[Feature] Add MiniCPM-SALA support (#30360)
Co-authored-by: Alex Nails <alex.nails@radixark.ai> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
cauphe
co-authored by
Alex Nails
Claude Opus 5
Alex Nails
Claude Opus 5
parent
d251fa2453
commit
092d85eb87
@@ -0,0 +1,49 @@
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.benchmark import marker
|
||||
from sglang.kernels.jit.minicpm_sala import get_block_table
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(
|
||||
est_time=20, stage="base-b-kernel-benchmark", runner_config="1-gpu-large"
|
||||
)
|
||||
|
||||
_HEAD_GROUP = 2
|
||||
_SPARSE_BLOCK_SIZE = 64
|
||||
_TOPK = 96
|
||||
|
||||
|
||||
def _make_valid_inputs(token_num: int, topk: int, device: str = "cuda"):
|
||||
"""Well-formed inputs shared by both expansion strategies.
|
||||
|
||||
``seqlen_q_max`` is tied to ``token_num`` so the per-token causal position
|
||||
(``token_pos_in_bs``) never indexes past ``block_table``.
|
||||
"""
|
||||
seqlen_q_max = token_num
|
||||
num_blocks = max(1, seqlen_q_max // _SPARSE_BLOCK_SIZE)
|
||||
torch.manual_seed(0)
|
||||
topk_idx = torch.randint(
|
||||
0, num_blocks, (_HEAD_GROUP, token_num, topk), dtype=torch.int32, device=device
|
||||
)
|
||||
block_table = torch.arange(
|
||||
1, seqlen_q_max + 1, dtype=torch.int32, device=device
|
||||
).reshape(1, seqlen_q_max)
|
||||
token_to_bs = torch.zeros((token_num,), dtype=torch.int32, device=device)
|
||||
token_pos_in_bs = torch.arange(1, token_num + 1, dtype=torch.int32, device=device)
|
||||
seqlen_q = torch.tensor([seqlen_q_max], dtype=torch.int32, device=device)
|
||||
return topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q
|
||||
|
||||
|
||||
@marker.parametrize("token_num", [2**n for n in range(9, 15)], [512, 4096])
|
||||
@marker.benchmark("provider", ["blockwise", "elementwise"])
|
||||
def benchmark(token_num: int, provider: str):
|
||||
inputs = _make_valid_inputs(token_num, _TOPK)
|
||||
|
||||
def fn(*args):
|
||||
return get_block_table(*args, elementwise=provider == "elementwise")
|
||||
|
||||
return marker.do_bench(fn, input_args=inputs)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
benchmark.run()
|
||||
@@ -0,0 +1,153 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.kernels.jit.minicpm_sala.get_block_table import get_block_table
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
_HEAD_GROUP = 2
|
||||
_SPARSE_BLOCK_SIZE = 64
|
||||
|
||||
|
||||
def _make_inputs(token_num, seqlen_q_max, topk, batch_size=1, device="cuda"):
|
||||
"""Build the same kind of inputs as the original CUDA kernel test."""
|
||||
topk_idx = torch.full(
|
||||
(_HEAD_GROUP, token_num, topk), -1, dtype=torch.int32, device=device
|
||||
)
|
||||
# Plant a few valid blocks at fixed positions, like the original UT.
|
||||
topk_idx[0, 32, 0:2] = torch.tensor([0, 1], dtype=torch.int32, device=device)
|
||||
topk_idx[1, 32, 0:2] = torch.tensor([0, 1], dtype=torch.int32, device=device)
|
||||
topk_idx[1, 64, 0:2] = torch.tensor([0, 1], dtype=torch.int32, device=device)
|
||||
topk_idx[0, 1000, 0:10] = torch.tensor(
|
||||
[0, 1, 5, 11, 14, 16, 17, 25, 26, 27], dtype=torch.int32, device=device
|
||||
)
|
||||
|
||||
block_table = torch.arange(
|
||||
1, seqlen_q_max * batch_size + 1, dtype=torch.int32, device=device
|
||||
).reshape(batch_size, seqlen_q_max)
|
||||
token_to_bs = torch.zeros((token_num,), dtype=torch.int32, device=device)
|
||||
token_pos_in_bs = torch.arange(1, token_num + 1, dtype=torch.int32, device=device)
|
||||
seqlen_q = torch.tensor([seqlen_q_max], dtype=torch.int32, device=device)
|
||||
return topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q
|
||||
|
||||
|
||||
def _make_valid_inputs(
|
||||
token_num,
|
||||
seqlen_q_max,
|
||||
topk,
|
||||
batch_size=1,
|
||||
head_group=_HEAD_GROUP,
|
||||
block_size=_SPARSE_BLOCK_SIZE,
|
||||
device="cuda",
|
||||
):
|
||||
"""Build inputs with only non-negative block indices."""
|
||||
num_blocks = seqlen_q_max // block_size
|
||||
torch.manual_seed(0)
|
||||
topk_idx = torch.randint(
|
||||
0, num_blocks, (head_group, token_num, topk), dtype=torch.int32, device=device
|
||||
)
|
||||
block_table = torch.arange(
|
||||
1, seqlen_q_max * batch_size + 1, dtype=torch.int32, device=device
|
||||
).reshape(batch_size, seqlen_q_max)
|
||||
token_to_bs = torch.zeros((token_num,), dtype=torch.int32, device=device)
|
||||
token_pos_in_bs = torch.arange(1, token_num + 1, dtype=torch.int32, device=device)
|
||||
seqlen_q = torch.tensor([seqlen_q_max], dtype=torch.int32, device=device)
|
||||
return topk_idx, block_table, token_to_bs, token_pos_in_bs, seqlen_q
|
||||
|
||||
|
||||
def _get_block_table_reference(
|
||||
topk_idx,
|
||||
block_table,
|
||||
token_to_bs,
|
||||
token_pos_in_bs,
|
||||
seqlen_q,
|
||||
block_size=_SPARSE_BLOCK_SIZE,
|
||||
):
|
||||
head_group = topk_idx.shape[0]
|
||||
token_num = topk_idx.shape[1]
|
||||
source = topk_idx.permute(1, 0, 2).unsqueeze(-1) * block_size + torch.arange(
|
||||
block_size, device=topk_idx.device
|
||||
)
|
||||
valid = (source >= 0) & (
|
||||
source
|
||||
< torch.minimum(seqlen_q[token_to_bs], token_pos_in_bs).view(token_num, 1, 1, 1)
|
||||
)
|
||||
gathered = torch.gather(
|
||||
block_table[token_to_bs],
|
||||
1,
|
||||
source.reshape(token_num, -1).clamp(0, block_table.shape[1] - 1),
|
||||
).view_as(source)
|
||||
heads = torch.arange(head_group, device=topk_idx.device).view(1, -1, 1, 1)
|
||||
return torch.where(valid, gathered * head_group + heads, 0).flatten(2)
|
||||
|
||||
|
||||
def test_get_block_table_supports_tp_local_head_group():
|
||||
inputs = _make_valid_inputs(64, 64, 96, head_group=1)
|
||||
expected = _get_block_table_reference(*inputs)
|
||||
actual = get_block_table(*inputs, head_group_num=1, elementwise=False)
|
||||
assert torch.equal(expected, actual)
|
||||
|
||||
|
||||
def _golden_check_blockwise(out_block_table, block_table, token_num):
|
||||
"""The assertions ported verbatim from the original kernel test."""
|
||||
# check token 32
|
||||
assert (out_block_table[32, 0] != 0).sum().item() == 33
|
||||
assert (out_block_table[32, 1] != 0).sum().item() == 33
|
||||
assert torch.equal(out_block_table[32, 0, 0:33], block_table[0][:33] * 2)
|
||||
assert torch.equal(out_block_table[32, 1, 0:33], block_table[0][:33] * 2 + 1)
|
||||
|
||||
# check token 64
|
||||
assert (out_block_table[64, 1] != 0).sum().item() == 65
|
||||
assert torch.equal(out_block_table[64, 1, 0:65], block_table[0][:65] * 2 + 1)
|
||||
|
||||
# check token 1000
|
||||
topk_blocks = [0, 1, 5, 11, 14, 16, 17, 25, 26, 27]
|
||||
tokens = []
|
||||
for b in topk_blocks:
|
||||
tokens.extend(range(b * _SPARSE_BLOCK_SIZE, (b + 1) * _SPARSE_BLOCK_SIZE))
|
||||
tokens = [t for t in tokens if t < token_num and t < 1001]
|
||||
assert (out_block_table[1000, 0] != 0).sum().item() == len(tokens)
|
||||
assert torch.equal(
|
||||
out_block_table[1000, 0, : len(tokens)], block_table[0][tokens] * 2
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("topk", [96, 128])
|
||||
def test_get_block_table_blockwise_golden(topk):
|
||||
token_num, seqlen_q_max = 8192, 8192
|
||||
inputs = _make_inputs(token_num, seqlen_q_max, topk)
|
||||
out = get_block_table(*inputs, elementwise=False)
|
||||
assert out.shape == (token_num, _HEAD_GROUP, topk * _SPARSE_BLOCK_SIZE)
|
||||
_golden_check_blockwise(out, inputs[1], token_num)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("topk", [96, 128])
|
||||
def test_get_block_table_strategies_match_reference(topk):
|
||||
"""Both expansion strategies match the Torch reference, including -1."""
|
||||
token_num, seqlen_q_max = 2048, 2048
|
||||
inputs = _make_inputs(token_num, seqlen_q_max, topk)
|
||||
expected = _get_block_table_reference(*inputs)
|
||||
assert torch.equal(expected, get_block_table(*inputs, elementwise=False))
|
||||
assert torch.equal(expected, get_block_table(*inputs, elementwise=True))
|
||||
|
||||
|
||||
@pytest.mark.parametrize(("topk", "block_size"), [(10, 32), (7, 128)])
|
||||
def test_get_block_table_supports_configured_layout(topk, block_size):
|
||||
token_num = seqlen_q_max = 256
|
||||
inputs = _make_valid_inputs(
|
||||
token_num,
|
||||
seqlen_q_max,
|
||||
topk,
|
||||
block_size=block_size,
|
||||
)
|
||||
expected = _get_block_table_reference(*inputs, block_size=block_size)
|
||||
kwargs = {"head_group_num": _HEAD_GROUP, "block_size": block_size}
|
||||
assert torch.equal(expected, get_block_table(*inputs, **kwargs, elementwise=False))
|
||||
assert torch.equal(expected, get_block_table(*inputs, **kwargs, elementwise=True))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,83 @@
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.minicpm.fuse_kernel import (
|
||||
fused_attn_pooling_online_topk_decode,
|
||||
)
|
||||
from sglang.srt.layers.attention.minicpm.sparse_utils import compress_k_core_new
|
||||
from sglang.test.ci.ci_register import register_cuda_ci
|
||||
|
||||
register_cuda_ci(est_time=10, stage="base-b-kernel-unit", runner_config="1-gpu-large")
|
||||
|
||||
|
||||
def test_compress_k_writes_each_head_once():
|
||||
"""Each compressed output must be produced once even with multiple KV heads."""
|
||||
key_cache = torch.arange(
|
||||
9 * 2 * 6,
|
||||
dtype=torch.float32,
|
||||
device="cuda",
|
||||
).reshape(9, 2, 6)
|
||||
original = key_cache.clone()
|
||||
token_table = torch.tensor([[0, 0, 0, 1, 2, 3]], dtype=torch.int32, device="cuda")
|
||||
compressed_table = torch.tensor([[6, 7, 8]], dtype=torch.int32, device="cuda")
|
||||
full_compressed = torch.empty((3, 2, 6), device="cuda")
|
||||
|
||||
compress_k_core_new(
|
||||
full_compressed,
|
||||
1,
|
||||
key_cache,
|
||||
token_table,
|
||||
compressed_table,
|
||||
torch.tensor([0, 4], dtype=torch.int32, device="cuda"),
|
||||
torch.tensor([1], dtype=torch.int32, device="cuda"),
|
||||
torch.tensor([0, 3], dtype=torch.int32, device="cuda"),
|
||||
2,
|
||||
2,
|
||||
6,
|
||||
)
|
||||
|
||||
expected = torch.stack(
|
||||
(
|
||||
original[6],
|
||||
original[0:2].mean(dim=0),
|
||||
original[2:4].mean(dim=0),
|
||||
)
|
||||
)
|
||||
torch.testing.assert_close(full_compressed, expected)
|
||||
torch.testing.assert_close(key_cache[7:9], expected[1:])
|
||||
|
||||
|
||||
def test_fused_decode_topk_skips_dense_rows():
|
||||
kernel = fused_attn_pooling_online_topk_decode(
|
||||
batch_size=2,
|
||||
groups=16,
|
||||
heads=16,
|
||||
dim=128,
|
||||
topk=8,
|
||||
pooled_k_len=8,
|
||||
dense_len=5,
|
||||
dtype_str="bfloat16",
|
||||
)
|
||||
topk_indices = torch.full((1, 2, 8), -1, dtype=torch.int32, device="cuda")
|
||||
topk_values = torch.full(
|
||||
(1, 2, 8), float("-inf"), dtype=torch.float32, device="cuda"
|
||||
)
|
||||
|
||||
kernel(
|
||||
torch.randn(32, 1, 128, dtype=torch.bfloat16, device="cuda"),
|
||||
torch.randn(4, 1, 128, dtype=torch.bfloat16, device="cuda"),
|
||||
torch.tensor([0, 1, 2], dtype=torch.int32, device="cuda"),
|
||||
torch.tensor([0, 2, 4], dtype=torch.int32, device="cuda"),
|
||||
torch.tensor([3, 7], dtype=torch.int32, device="cuda"),
|
||||
topk_indices,
|
||||
topk_values,
|
||||
)
|
||||
|
||||
assert torch.all(topk_indices[:, 0] == -1)
|
||||
assert torch.any(topk_indices[:, 1] >= 0)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
import sys
|
||||
|
||||
sys.exit(pytest.main([__file__, "-v", "-s"]))
|
||||
@@ -0,0 +1,349 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.configs.hybrid_arch import (
|
||||
hybrid_lightning_config,
|
||||
mambaish_config,
|
||||
)
|
||||
from sglang.srt.configs.linear_attn_model_registry import (
|
||||
get_linear_attn_config,
|
||||
get_linear_attn_spec_by_arch,
|
||||
)
|
||||
from sglang.srt.configs.mamba_utils import Mamba2CacheParams
|
||||
from sglang.srt.configs.minicpm import MiniCPMHybridConfig
|
||||
from sglang.srt.layers.attention.hybrid_linear_attn_backend import (
|
||||
MambaAttnBackendBase,
|
||||
)
|
||||
from sglang.srt.layers.attention.linear.lightning_backend import (
|
||||
LightningAttentionBackend,
|
||||
)
|
||||
from sglang.srt.models import minicpm as minicpm_module
|
||||
from sglang.srt.models.minicpm import (
|
||||
MiniCPMAttention,
|
||||
MiniCPMDecoderLayer,
|
||||
MiniCPMLightningMixer,
|
||||
)
|
||||
from sglang.srt.runtime_context import get_parallel
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def test_minicpm_lightning_config_defaults_are_complete():
|
||||
"""A checkpoint missing optional SALA fields must still define every model input."""
|
||||
config = MiniCPMHybridConfig()
|
||||
|
||||
assert config.scale_emb == 12
|
||||
assert config.scale_depth == 1.4
|
||||
assert config.dim_model_base == 256
|
||||
assert config.lightning_use_rope is True
|
||||
assert config.use_output_gate is False
|
||||
assert config.attention_bias is False
|
||||
assert config.use_output_norm is False
|
||||
assert config.qk_norm is True
|
||||
|
||||
|
||||
def test_minicpm_empty_mixer_types_default_to_full_attention():
|
||||
config = MiniCPMHybridConfig(num_hidden_layers=3, mixer_types=[])
|
||||
|
||||
assert config.mixer_types == ["minicpm4", "minicpm4", "minicpm4"]
|
||||
assert config.full_attention_layer_ids == [0, 1, 2]
|
||||
|
||||
|
||||
def test_minicpm_sparse_config_uses_nested_fields_only():
|
||||
sparse_config = {
|
||||
"block_size": 64,
|
||||
"dense_len": 8192,
|
||||
"init_blocks": 1,
|
||||
"kernel_size": 32,
|
||||
"kernel_stride": 16,
|
||||
"topk": 64,
|
||||
"window_size": 2048,
|
||||
}
|
||||
config = MiniCPMHybridConfig(sparse_config=sparse_config)
|
||||
|
||||
assert config.has_minicpm_sparse_attention
|
||||
assert config.sparse_config == sparse_config
|
||||
assert not hasattr(config, "sparse_dense_len")
|
||||
|
||||
|
||||
def test_minicpm_short_mixer_pattern_repeats_to_layer_count():
|
||||
config = MiniCPMHybridConfig(
|
||||
num_hidden_layers=5,
|
||||
mixer_types=["minicpm4", "lightning-attn"],
|
||||
lightning_nkv=32,
|
||||
)
|
||||
|
||||
assert config.mixer_types == [
|
||||
"minicpm4",
|
||||
"lightning-attn",
|
||||
"minicpm4",
|
||||
"lightning-attn",
|
||||
"minicpm4",
|
||||
]
|
||||
assert config.full_attention_layer_ids == [0, 2, 4]
|
||||
assert config.lightning_layer_ids == [1, 3]
|
||||
|
||||
|
||||
def test_minicpm_mixer_aliases_are_canonicalized():
|
||||
config = MiniCPMHybridConfig(
|
||||
num_hidden_layers=4,
|
||||
mixer_types=["attention", "lightning_attn"],
|
||||
lightning_nkv=32,
|
||||
)
|
||||
|
||||
assert config.mixer_types == [
|
||||
"minicpm4",
|
||||
"lightning-attn",
|
||||
"minicpm4",
|
||||
"lightning-attn",
|
||||
]
|
||||
|
||||
|
||||
def test_minicpm_rejects_more_mixer_types_than_layers():
|
||||
with pytest.raises(ValueError, match="Invalid number of mixer types: 3"):
|
||||
MiniCPMHybridConfig(
|
||||
num_hidden_layers=2,
|
||||
mixer_types=["minicpm4", "lightning", "minicpm4"],
|
||||
)
|
||||
|
||||
|
||||
def test_minicpm_lightning_dimensions_fall_back_to_base_attention():
|
||||
config = MiniCPMHybridConfig(
|
||||
hidden_size=96,
|
||||
num_attention_heads=6,
|
||||
num_key_value_heads=3,
|
||||
head_dim=None,
|
||||
lightning_nh=None,
|
||||
lightning_nkv=None,
|
||||
lightning_head_dim=None,
|
||||
)
|
||||
|
||||
assert config.head_dim == 16
|
||||
assert config.lightning_nh == 6
|
||||
assert config.lightning_nkv == 3
|
||||
assert config.lightning_head_dim == 16
|
||||
|
||||
|
||||
def test_minicpm_rejects_lightning_gqa():
|
||||
with pytest.raises(ValueError, match="seg_la backend does not support GQA"):
|
||||
MiniCPMHybridConfig(
|
||||
num_attention_heads=6,
|
||||
num_key_value_heads=3,
|
||||
mixer_types=["lightning-attn"],
|
||||
)
|
||||
|
||||
|
||||
def test_minicpm_lightning_idle_batch_returns_empty_output():
|
||||
"""An idle DP rank must return empty output instead of reducing empty tensors."""
|
||||
mixer = MiniCPMLightningMixer.__new__(MiniCPMLightningMixer)
|
||||
torch.nn.Module.__init__(mixer)
|
||||
mixer.hidden_size = 8
|
||||
forward_batch = SimpleNamespace(forward_mode=SimpleNamespace(is_idle=lambda: True))
|
||||
|
||||
output = mixer.forward(
|
||||
positions=torch.empty(0, dtype=torch.int64),
|
||||
hidden_states=torch.empty(0, 4),
|
||||
forward_batch=forward_batch,
|
||||
)
|
||||
|
||||
assert output.shape == (0, 8)
|
||||
|
||||
|
||||
def test_minicpm_lightning_attention_bias_applies_to_every_projection():
|
||||
"""Enabling attention bias must cover every Lightning projection."""
|
||||
with get_parallel().override(tp_size=1, tp_rank=0):
|
||||
mixer = MiniCPMLightningMixer(
|
||||
hidden_size=8,
|
||||
num_heads=2,
|
||||
num_kv_heads=2,
|
||||
head_dim=4,
|
||||
use_rope=False,
|
||||
use_output_gate=True,
|
||||
attention_bias=True,
|
||||
qk_norm=False,
|
||||
)
|
||||
|
||||
assert mixer.qkv_proj.bias is not None
|
||||
assert mixer.o_proj.bias is not None
|
||||
assert mixer.z_proj.bias is not None
|
||||
|
||||
|
||||
def test_minicpm_lightning_rejects_unknown_scale():
|
||||
with (
|
||||
get_parallel().override(tp_size=1, tp_rank=0),
|
||||
pytest.raises(ValueError, match="Unsupported lightning scale"),
|
||||
):
|
||||
MiniCPMLightningMixer(
|
||||
hidden_size=8,
|
||||
num_heads=2,
|
||||
num_kv_heads=2,
|
||||
head_dim=4,
|
||||
use_rope=False,
|
||||
qk_norm=False,
|
||||
scale="unknown",
|
||||
)
|
||||
|
||||
|
||||
def test_minicpm_full_attention_bias_applies_to_every_projection():
|
||||
"""Enabling attention bias must cover every full-attention projection."""
|
||||
with get_parallel().override(tp_size=1, tp_rank=0):
|
||||
mixer = MiniCPMAttention(
|
||||
hidden_size=8,
|
||||
num_heads=2,
|
||||
num_kv_heads=2,
|
||||
attn_use_rope=False,
|
||||
use_output_gate=True,
|
||||
attention_bias=True,
|
||||
)
|
||||
|
||||
assert mixer.qkv_proj.bias is not None
|
||||
assert mixer.o_proj.bias is not None
|
||||
assert mixer.o_gate.bias is not None
|
||||
|
||||
|
||||
def test_minicpm_full_attention_uses_configured_head_dim(monkeypatch):
|
||||
monkeypatch.setattr(minicpm_module, "SiluAndMul", torch.nn.Identity)
|
||||
config = MiniCPMHybridConfig(
|
||||
hidden_size=16,
|
||||
num_hidden_layers=1,
|
||||
num_attention_heads=2,
|
||||
num_key_value_heads=2,
|
||||
head_dim=6,
|
||||
intermediate_size=32,
|
||||
attn_use_rope=False,
|
||||
)
|
||||
|
||||
with get_parallel().override(tp_size=1, tp_rank=0):
|
||||
layer = MiniCPMDecoderLayer(config)
|
||||
|
||||
assert layer.self_attn.head_dim == 6
|
||||
assert layer.self_attn.q_size == 12
|
||||
assert layer.self_attn.kv_size == 12
|
||||
|
||||
|
||||
def test_minicpm_lightning_reuses_shared_backend_and_cache_shape():
|
||||
config = MiniCPMHybridConfig(
|
||||
num_hidden_layers=2,
|
||||
mixer_types=["lightning", "minicpm4"],
|
||||
lightning_nh=4,
|
||||
lightning_nkv=4,
|
||||
lightning_head_dim=64,
|
||||
)
|
||||
|
||||
model_config = SimpleNamespace(
|
||||
hf_config=config,
|
||||
linear_attn_registry_result=get_linear_attn_config(config),
|
||||
)
|
||||
assert hybrid_lightning_config(model_config) is config
|
||||
assert mambaish_config(model_config) is config
|
||||
|
||||
with get_parallel().override(attn_tp_size=1):
|
||||
cache = config.mamba2_cache_params
|
||||
assert isinstance(cache, Mamba2CacheParams)
|
||||
assert cache.layers == [0]
|
||||
assert cache.shape.conv == [(0, 0)]
|
||||
assert cache.shape.temporal == (4, 64, 64)
|
||||
assert config.num_linear_key_value_heads == 4
|
||||
|
||||
with get_parallel().override(attn_tp_size=1, attn_tp_rank=0):
|
||||
slopes = LightningAttentionBackend._build_slope_tensor(
|
||||
4, 2, device="cpu", layerwise_decay=False
|
||||
)
|
||||
assert len(slopes) == 2
|
||||
assert slopes[0].equal(slopes[1])
|
||||
|
||||
|
||||
def test_non_lightning_minicpm_is_not_classified_as_linear_attention():
|
||||
config = MiniCPMHybridConfig(
|
||||
num_hidden_layers=1,
|
||||
mixer_types=["minicpm4"],
|
||||
sparse_config={},
|
||||
)
|
||||
model_config = SimpleNamespace(
|
||||
hf_config=config,
|
||||
linear_attn_registry_result=get_linear_attn_config(config),
|
||||
)
|
||||
|
||||
assert hybrid_lightning_config(model_config) is None
|
||||
assert mambaish_config(model_config) is None
|
||||
for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
|
||||
assert get_linear_attn_spec_by_arch(architecture) is None
|
||||
|
||||
|
||||
def test_lightning_backend_reads_structural_linear_config(monkeypatch):
|
||||
def fake_base_init(self, model_runner):
|
||||
self.topk = 1
|
||||
|
||||
monkeypatch.setattr(MambaAttnBackendBase, "__init__", fake_base_init)
|
||||
config = SimpleNamespace(
|
||||
num_attention_heads=8,
|
||||
num_linear_key_value_heads=4,
|
||||
num_hidden_layers=2,
|
||||
lightning_layerwise_decay=False,
|
||||
)
|
||||
model_runner = SimpleNamespace(
|
||||
req_to_token_pool=SimpleNamespace(
|
||||
mamba_pool=SimpleNamespace(
|
||||
mamba_cache=SimpleNamespace(conv=[torch.empty(0)])
|
||||
)
|
||||
),
|
||||
sliding_window_size=None,
|
||||
model_config=SimpleNamespace(
|
||||
hf_config=config,
|
||||
is_encoder_decoder=False,
|
||||
context_len=128,
|
||||
block=256,
|
||||
),
|
||||
device="cpu",
|
||||
kv_cache_dtype=torch.float32,
|
||||
kv_cache_dtype_str="float32",
|
||||
)
|
||||
|
||||
with get_parallel().override(attn_tp_size=1, attn_tp_rank=0):
|
||||
backend = LightningAttentionBackend(model_runner)
|
||||
|
||||
assert [slope.shape for slope in backend.tp_slope] == [(4, 1, 1), (4, 1, 1)]
|
||||
assert backend.tp_slope[0].equal(backend.tp_slope[1])
|
||||
|
||||
|
||||
def test_lightning_backend_uses_layer_scale(monkeypatch):
|
||||
"""Each layer's attention scale must reach the linear-attention computation."""
|
||||
captured = {}
|
||||
|
||||
def fake_seg_la_fwd(**kwargs):
|
||||
captured.update(kwargs)
|
||||
return kwargs["q"]
|
||||
|
||||
monkeypatch.setattr(
|
||||
"sglang.srt.layers.attention.linear.lightning_backend.seg_la_fwd",
|
||||
fake_seg_la_fwd,
|
||||
)
|
||||
backend = LightningAttentionBackend.__new__(LightningAttentionBackend)
|
||||
backend.tp_slope = [torch.ones(1, 1, 1)]
|
||||
layer = SimpleNamespace(layer_id=0, scaling=0.25)
|
||||
metadata = SimpleNamespace(
|
||||
batch_size=1,
|
||||
query_start_loc=torch.tensor([0, 1]),
|
||||
has_initial_states=torch.tensor([False]),
|
||||
)
|
||||
q = torch.ones(1, 1, 1)
|
||||
|
||||
backend._linear_attention_entry(
|
||||
q=q,
|
||||
k=q,
|
||||
v=q,
|
||||
kv_cache=torch.zeros(1, 1, 1, 1),
|
||||
state_indices_tensor=torch.tensor([0]),
|
||||
metadata=metadata,
|
||||
layer=layer,
|
||||
)
|
||||
|
||||
assert captured["softmax_scale"] == 0.25
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,55 @@
|
||||
import sys
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.disaggregation.decode import (
|
||||
DecodeReqToTokenPool,
|
||||
HybridMambaDecodeReqToTokenPool,
|
||||
)
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=1, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _init_decode_pool(pool):
|
||||
DecodeReqToTokenPool.__init__(
|
||||
pool,
|
||||
size=1,
|
||||
max_context_len=4,
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
pre_alloc_size=1,
|
||||
)
|
||||
return pool
|
||||
|
||||
|
||||
def test_decode_pool_reports_physical_capacity():
|
||||
pool = _init_decode_pool(DecodeReqToTokenPool.__new__(DecodeReqToTokenPool))
|
||||
|
||||
assert pool.schedulable_token_capacity(17) == 17
|
||||
|
||||
|
||||
def test_decode_pool_supports_noop_aux_cache_contract():
|
||||
pool = _init_decode_pool(DecodeReqToTokenPool.__new__(DecodeReqToTokenPool))
|
||||
req_to_token = pool.req_to_token.clone()
|
||||
|
||||
pool.alloc_aux_to_lengths(
|
||||
req_pool_indices_cpu=torch.tensor([1]),
|
||||
target_seq_lens_cpu=torch.tensor([3]),
|
||||
)
|
||||
pool.reset_aux_cache_allocator()
|
||||
|
||||
assert torch.equal(pool.req_to_token, req_to_token)
|
||||
|
||||
|
||||
def test_hybrid_decode_pool_initializes_aux_cache_contract():
|
||||
pool = _init_decode_pool(
|
||||
HybridMambaDecodeReqToTokenPool.__new__(HybridMambaDecodeReqToTokenPool)
|
||||
)
|
||||
|
||||
assert pool.schedulable_token_capacity(17) == 17
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
@@ -0,0 +1,91 @@
|
||||
import sys
|
||||
import unittest
|
||||
from types import SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
module: MagicMock()
|
||||
for module in (
|
||||
"sgl_kernel",
|
||||
"sgl_kernel.quantization",
|
||||
"sgl_kernel.scalar_type",
|
||||
)
|
||||
},
|
||||
):
|
||||
from sglang.srt.layers.attention.flashattention_backend import (
|
||||
FlashAttentionBackend,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _backend():
|
||||
backend = FlashAttentionBackend.__new__(FlashAttentionBackend)
|
||||
backend.page_size = 1
|
||||
backend.kv_cache_dtype = torch.float16
|
||||
backend.kv_cache_dtype_str = "float8_e4m3fn"
|
||||
backend.kv_cache_is_mxfp8 = False
|
||||
backend.fa_impl_ver = 3
|
||||
backend.num_splits = 4
|
||||
return backend
|
||||
|
||||
|
||||
class TestFlashAttentionPagedMHA(unittest.TestCase):
|
||||
def test_get_paged_mha_kv_cache_supports_head_groups(self):
|
||||
backend = _backend()
|
||||
backend.token_to_kv_pool = SimpleNamespace(
|
||||
get_kv_buffer=Mock(
|
||||
return_value=(
|
||||
torch.empty(8, 2, 16),
|
||||
torch.empty(8, 2, 16),
|
||||
)
|
||||
)
|
||||
)
|
||||
layer = SimpleNamespace(
|
||||
layer_id=3,
|
||||
tp_k_head_num=2,
|
||||
tp_v_head_num=2,
|
||||
head_dim=16,
|
||||
v_head_dim=16,
|
||||
)
|
||||
|
||||
key_cache, value_cache = backend.get_paged_mha_kv_cache(
|
||||
layer,
|
||||
head_group_num=2,
|
||||
)
|
||||
|
||||
self.assertEqual(key_cache.shape, (16, 1, 1, 16))
|
||||
self.assertEqual(value_cache.shape, (16, 1, 1, 16))
|
||||
|
||||
def test_prepare_paged_mha_query_reuses_fa_scaling_policy(self):
|
||||
backend = _backend()
|
||||
layer = SimpleNamespace(
|
||||
head_dim=16,
|
||||
k_scale=torch.tensor(2.0),
|
||||
v_scale=torch.tensor(4.0),
|
||||
)
|
||||
q = torch.ones(2, 16, dtype=torch.bfloat16)
|
||||
|
||||
q, _, _, k_descale, v_descale = backend.prepare_paged_mha_query(
|
||||
q,
|
||||
None,
|
||||
None,
|
||||
layer,
|
||||
logical_batch_size=2,
|
||||
kv_head_num=1,
|
||||
is_prefill=True,
|
||||
)
|
||||
|
||||
self.assertEqual(q.dtype, torch.float16)
|
||||
self.assertEqual(k_descale.shape, (2, 1))
|
||||
self.assertEqual(v_descale.shape, (2, 1))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,214 @@
|
||||
import sys
|
||||
import unittest
|
||||
from types import ModuleType, SimpleNamespace
|
||||
from unittest.mock import MagicMock, Mock, patch
|
||||
|
||||
import torch
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
with patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
module: MagicMock()
|
||||
for module in (
|
||||
"sgl_kernel",
|
||||
"sgl_kernel.quantization",
|
||||
"sgl_kernel.scalar_type",
|
||||
)
|
||||
},
|
||||
):
|
||||
from sglang.srt.layers.attention.minicpm import attention_adapter as adapter_module
|
||||
from sglang.srt.layers.attention.minicpm.attention_adapter import (
|
||||
MiniCPMFlashAttentionAdapter,
|
||||
MiniCPMFlashInferAdapter,
|
||||
)
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
def _metadata(rows=1):
|
||||
return SimpleNamespace(
|
||||
sparse_page_table=torch.zeros((rows, 4), dtype=torch.int32),
|
||||
sparse_cache_seqlens_int32=torch.full(
|
||||
(rows,),
|
||||
4,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
sparse_cu_seqlens_q=torch.arange(rows + 1, dtype=torch.int32),
|
||||
sparse_cu_seqlens_k=torch.arange(
|
||||
0,
|
||||
(rows + 1) * 4,
|
||||
4,
|
||||
dtype=torch.int32,
|
||||
),
|
||||
sparse_max_seq_len_q=1,
|
||||
max_seq_len_q=1,
|
||||
)
|
||||
|
||||
|
||||
class TestMiniCPMAttentionAdapter(unittest.TestCase):
|
||||
def test_flashattention_adapter_owns_kernel_arguments(self):
|
||||
expected = torch.ones(1, 1, 1)
|
||||
flash_attn_backend = SimpleNamespace(
|
||||
num_splits=4,
|
||||
fa_impl_ver=3,
|
||||
)
|
||||
adapter = MiniCPMFlashAttentionAdapter(flash_attn_backend)
|
||||
metadata = _metadata()
|
||||
layer = SimpleNamespace(scaling=0.125, logit_cap=0.0)
|
||||
k_descale = torch.tensor([[2.0]])
|
||||
v_descale = torch.tensor([[4.0]])
|
||||
|
||||
with patch.object(
|
||||
adapter_module,
|
||||
"flash_attn_with_kvcache",
|
||||
return_value=expected,
|
||||
) as kernel:
|
||||
result = adapter.forward(
|
||||
torch.ones(1, 1, 1),
|
||||
torch.ones(4, 1, 1, 1),
|
||||
torch.ones(4, 1, 1, 1),
|
||||
metadata,
|
||||
layer,
|
||||
is_prefill=True,
|
||||
k_descale=k_descale,
|
||||
v_descale=v_descale,
|
||||
)
|
||||
|
||||
self.assertIs(result, expected)
|
||||
kwargs = kernel.call_args.kwargs
|
||||
self.assertIs(kwargs["page_table"], metadata.sparse_page_table)
|
||||
self.assertIs(kwargs["k_descale"], k_descale)
|
||||
self.assertIs(kwargs["v_descale"], v_descale)
|
||||
self.assertEqual(kwargs["num_splits"], 4)
|
||||
self.assertEqual(kwargs["ver"], 3)
|
||||
|
||||
def test_flashinfer_prefill_plans_once_and_executes_each_layer(self):
|
||||
adapter = MiniCPMFlashInferAdapter.__new__(MiniCPMFlashInferAdapter)
|
||||
adapter.prefill_planned = False
|
||||
adapter._prepare = Mock()
|
||||
adapter.active_rows = torch.tensor([0], dtype=torch.int32)
|
||||
adapter.active_kv_indptr = torch.tensor([0, 4], dtype=torch.int32)
|
||||
adapter.active_kv_indices = torch.empty(4, dtype=torch.int32)
|
||||
expected = torch.ones(1, 1, 1)
|
||||
adapter.active_wrapper = SimpleNamespace(forward=Mock(return_value=expected))
|
||||
metadata = _metadata()
|
||||
layer = SimpleNamespace(
|
||||
scaling=0.125,
|
||||
logit_cap=0.0,
|
||||
k_scale_float=1.0,
|
||||
v_scale_float=1.0,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
adapter_module,
|
||||
"create_flashinfer_kv_indices_triton",
|
||||
) as index_kernel:
|
||||
first = adapter.forward(
|
||||
torch.ones(1, 1, 1),
|
||||
torch.ones(4, 1, 1, 1),
|
||||
torch.ones(4, 1, 1, 1),
|
||||
metadata,
|
||||
layer,
|
||||
is_prefill=True,
|
||||
)
|
||||
second = adapter.forward(
|
||||
torch.ones(1, 1, 1),
|
||||
torch.ones(4, 1, 1, 1),
|
||||
torch.ones(4, 1, 1, 1),
|
||||
metadata,
|
||||
layer,
|
||||
is_prefill=True,
|
||||
)
|
||||
|
||||
self.assertIs(first, expected)
|
||||
self.assertIs(second, expected)
|
||||
adapter._prepare.assert_called_once_with(metadata, is_prefill=True)
|
||||
self.assertEqual(index_kernel.__getitem__.return_value.call_count, 2)
|
||||
self.assertEqual(adapter.active_wrapper.forward.call_count, 2)
|
||||
|
||||
def test_flashinfer_graph_uses_backend_wrapper_cache(self):
|
||||
adapter = MiniCPMFlashInferAdapter.__new__(MiniCPMFlashInferAdapter)
|
||||
adapter.device = torch.device("cpu")
|
||||
adapter.head_group_num = 2
|
||||
adapter.num_qo_heads = 4
|
||||
adapter.num_kv_heads = 1
|
||||
adapter.head_dim = 16
|
||||
adapter.page_size = 1
|
||||
adapter.max_kv_tokens_per_row = 4
|
||||
adapter.q_dtype = torch.float16
|
||||
adapter.kv_dtype = torch.float16
|
||||
adapter.kv_indptr = torch.zeros(3, dtype=torch.int32)
|
||||
adapter.kv_indices = torch.zeros(8, dtype=torch.int32)
|
||||
adapter.kv_last_page_len = torch.ones(2, dtype=torch.int32)
|
||||
adapter.rows = torch.arange(2, dtype=torch.int32)
|
||||
wrapper = SimpleNamespace(begin_forward=Mock())
|
||||
adapter.flashinfer_backend = SimpleNamespace(
|
||||
get_cuda_graph_decode_wrappers=Mock(return_value=[wrapper]),
|
||||
)
|
||||
metadata = _metadata(rows=2)
|
||||
|
||||
adapter.prepare_forward(
|
||||
metadata,
|
||||
is_prefill=False,
|
||||
graph=True,
|
||||
)
|
||||
|
||||
adapter.flashinfer_backend.get_cuda_graph_decode_wrappers.assert_called_once_with(
|
||||
bs=1,
|
||||
num_tokens=2,
|
||||
)
|
||||
wrapper.begin_forward.assert_called_once()
|
||||
self.assertIs(adapter.active_wrapper, wrapper)
|
||||
|
||||
def test_flashinfer_decode_indices_cover_dense_rows(self):
|
||||
wrapper = SimpleNamespace(begin_forward=Mock())
|
||||
flashinfer_backend = SimpleNamespace(decode_wrappers=[wrapper])
|
||||
flashinfer_backend_module = ModuleType(
|
||||
"sglang.srt.layers.attention.flashinfer_backend"
|
||||
)
|
||||
flashinfer_backend_module.FlashInferAttnBackend = Mock(
|
||||
return_value=flashinfer_backend
|
||||
)
|
||||
model_runner = SimpleNamespace(
|
||||
device=torch.device("cpu"),
|
||||
dtype=torch.float16,
|
||||
kv_cache_dtype=torch.float16,
|
||||
req_to_token_pool=SimpleNamespace(size=1),
|
||||
)
|
||||
|
||||
with (
|
||||
patch.object(adapter_module, "is_flashinfer_available", return_value=True),
|
||||
patch.dict(
|
||||
sys.modules,
|
||||
{
|
||||
"sglang.srt.layers.attention.flashinfer_backend": (
|
||||
flashinfer_backend_module
|
||||
),
|
||||
},
|
||||
),
|
||||
):
|
||||
adapter = MiniCPMFlashInferAdapter(
|
||||
model_runner,
|
||||
head_group_num=2,
|
||||
heads_per_group=16,
|
||||
head_dim=128,
|
||||
page_size=1,
|
||||
max_kv_tokens_per_row=7,
|
||||
)
|
||||
|
||||
metadata = _metadata(rows=2)
|
||||
metadata.sparse_cache_seqlens_int32.fill_(7)
|
||||
adapter.prepare_forward(
|
||||
metadata,
|
||||
is_prefill=False,
|
||||
graph=False,
|
||||
)
|
||||
|
||||
self.assertEqual(adapter.kv_indices.numel(), 14)
|
||||
self.assertEqual(adapter.active_kv_indices.numel(), 14)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -0,0 +1,344 @@
|
||||
import sys
|
||||
from types import SimpleNamespace
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from sglang.srt.layers.attention.minicpm.cache import (
|
||||
attach_compressed_cache,
|
||||
)
|
||||
from sglang.srt.managers.scheduler_components.invariant_checker import (
|
||||
SchedulerInvariantChecker,
|
||||
)
|
||||
from sglang.srt.managers.scheduler_components.pool_stats_observer import (
|
||||
SchedulerPoolStatsObserver,
|
||||
)
|
||||
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
|
||||
from sglang.srt.session.streaming_session import SessionSlot, StreamingSession
|
||||
from sglang.test.ci.ci_register import register_cpu_ci
|
||||
|
||||
register_cpu_ci(est_time=5, suite="base-a-test-cpu")
|
||||
|
||||
|
||||
class RecordingAllocator:
|
||||
def __init__(self, capacity: int):
|
||||
self.capacity = capacity
|
||||
self.page_size = 1
|
||||
self.next_slot = 1
|
||||
self.live: set[int] = set()
|
||||
|
||||
@property
|
||||
def size(self):
|
||||
return self.capacity
|
||||
|
||||
def alloc(self, size: int):
|
||||
if size > self.available_size():
|
||||
return None
|
||||
slots = torch.arange(self.next_slot, self.next_slot + size, dtype=torch.int64)
|
||||
self.next_slot += size
|
||||
self.live.update(slots.tolist())
|
||||
return slots
|
||||
|
||||
def free(self, slots: torch.Tensor):
|
||||
self.live.difference_update(slots.tolist())
|
||||
|
||||
def available_size(self):
|
||||
return self.capacity - len(self.live)
|
||||
|
||||
def clear(self):
|
||||
self.next_slot = 1
|
||||
self.live.clear()
|
||||
|
||||
|
||||
def make_pool_and_req(capacity: int = 64):
|
||||
allocator = RecordingAllocator(capacity)
|
||||
pool = ReqToTokenPool(
|
||||
size=2,
|
||||
max_context_len=64,
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
attach_compressed_cache(
|
||||
pool,
|
||||
allocator,
|
||||
kernel_size=4,
|
||||
kernel_stride=2,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
req = SimpleNamespace(
|
||||
req_pool_idx=None,
|
||||
inflight_middle_chunks=0,
|
||||
kv_committed_len=0,
|
||||
)
|
||||
req_pool_idx = pool.alloc([req])[0]
|
||||
return pool, req, req_pool_idx, allocator
|
||||
|
||||
|
||||
def alloc_extend(pool, req_pool_idx: int, seq_len: int):
|
||||
pool.alloc_aux_to_lengths(
|
||||
req_pool_indices_cpu=torch.tensor([req_pool_idx], dtype=torch.int64),
|
||||
target_seq_lens_cpu=torch.tensor([seq_len], dtype=torch.int64),
|
||||
)
|
||||
|
||||
|
||||
def test_extend_allocates_at_sparse_boundaries():
|
||||
pool, _, req_pool_idx, allocator = make_pool_and_req()
|
||||
cache = pool._aux_cache
|
||||
|
||||
alloc_extend(pool, req_pool_idx, seq_len=3)
|
||||
assert allocator.available_size() == 39
|
||||
assert len(cache.free_slots) == 25
|
||||
|
||||
alloc_extend(pool, req_pool_idx, seq_len=4)
|
||||
assert allocator.available_size() == 39
|
||||
assert len(cache.free_slots) == 24
|
||||
|
||||
alloc_extend(pool, req_pool_idx, seq_len=16)
|
||||
assert allocator.available_size() == 39
|
||||
assert len(cache.free_slots) == 17
|
||||
|
||||
|
||||
def test_chunk_reuse_only_allocates_new_sparse_slots():
|
||||
pool, _, req_pool_idx, _ = make_pool_and_req()
|
||||
cache = pool._aux_cache
|
||||
|
||||
alloc_extend(pool, req_pool_idx, seq_len=8)
|
||||
assert len(cache.free_slots) == 22
|
||||
|
||||
alloc_extend(pool, req_pool_idx, seq_len=12)
|
||||
assert len(cache.free_slots) == 20
|
||||
|
||||
alloc_extend(pool, req_pool_idx, seq_len=12)
|
||||
assert len(cache.free_slots) == 20
|
||||
|
||||
|
||||
def test_decode_does_not_duplicate_sparse_slots():
|
||||
"""Retrying the same decode position must not allocate duplicate cache slots."""
|
||||
pool, _, req_pool_idx, _ = make_pool_and_req()
|
||||
cache = pool._aux_cache
|
||||
alloc_extend(pool, req_pool_idx, seq_len=15)
|
||||
|
||||
pool.alloc_aux_to_lengths(
|
||||
req_pool_indices_cpu=torch.tensor([req_pool_idx], dtype=torch.int64),
|
||||
target_seq_lens_cpu=torch.tensor([16], dtype=torch.int64),
|
||||
)
|
||||
available_after_first_decode = len(cache.free_slots)
|
||||
|
||||
pool.alloc_aux_to_lengths(
|
||||
req_pool_indices_cpu=torch.tensor([req_pool_idx], dtype=torch.int64),
|
||||
target_seq_lens_cpu=torch.tensor([16], dtype=torch.int64),
|
||||
)
|
||||
assert available_after_first_decode == 17
|
||||
assert len(cache.free_slots) == available_after_first_decode
|
||||
|
||||
|
||||
def test_reserve_leaves_only_dense_capacity_visible():
|
||||
allocator = RecordingAllocator(capacity=69)
|
||||
pool = ReqToTokenPool(
|
||||
size=2,
|
||||
max_context_len=64,
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
attach_compressed_cache(
|
||||
pool,
|
||||
allocator,
|
||||
kernel_size=32,
|
||||
kernel_stride=16,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
|
||||
assert allocator.available_size() == 64
|
||||
assert len(pool._aux_cache.reserved_slots) == 5
|
||||
assert pool.schedulable_token_capacity(69) == 64
|
||||
|
||||
|
||||
def test_reserved_slots_are_excluded_from_full_pool_invariant():
|
||||
pool, _, _, allocator = make_pool_and_req(capacity=69)
|
||||
checker = SchedulerInvariantChecker(
|
||||
is_hybrid_swa=False,
|
||||
is_hybrid_ssm=True,
|
||||
disaggregation_mode=None,
|
||||
page_size=1,
|
||||
full_tokens_per_layer=None,
|
||||
swa_tokens_per_layer=None,
|
||||
max_total_num_tokens=64,
|
||||
tree_cache=SimpleNamespace(
|
||||
supports_mamba=lambda: False,
|
||||
protected_size=lambda: 0,
|
||||
),
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
req_to_token_pool=pool,
|
||||
pool_stats_observer=SimpleNamespace(session_held_tokens=lambda: 0),
|
||||
get_last_batch=lambda: None,
|
||||
get_running_batch=lambda: None,
|
||||
)
|
||||
|
||||
leak, message = checker._check_full_pool(
|
||||
SimpleNamespace(
|
||||
full_available_size=allocator.available_size(), full_evictable_size=0
|
||||
)
|
||||
)
|
||||
|
||||
assert not leak, message
|
||||
|
||||
|
||||
def test_hybrid_pool_stats_exclude_reserved_slots():
|
||||
pool, _, _, allocator = make_pool_and_req(capacity=69)
|
||||
pool.mamba_allocator = SimpleNamespace(available_size=lambda: 1)
|
||||
pool.mamba_pool = SimpleNamespace(size=1)
|
||||
observer = SchedulerPoolStatsObserver(
|
||||
tree_cache=SimpleNamespace(supports_mamba=lambda: False),
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
req_to_token_pool=pool,
|
||||
session_controller=None,
|
||||
hisparse_coordinator=None,
|
||||
is_hybrid_swa=False,
|
||||
is_hybrid_ssm=True,
|
||||
enable_hisparse=False,
|
||||
full_tokens_per_layer=None,
|
||||
swa_tokens_per_layer=None,
|
||||
max_total_num_tokens=42,
|
||||
get_last_batch=lambda: None,
|
||||
get_running_batch=lambda: None,
|
||||
)
|
||||
|
||||
stats = observer._get_mamba_token_info()
|
||||
|
||||
assert stats.full_num_used == 0
|
||||
assert stats.full_token_usage == 0
|
||||
|
||||
|
||||
def test_streaming_session_release_frees_compressed_slots():
|
||||
pool, _, req_pool_idx, allocator = make_pool_and_req()
|
||||
alloc_extend(pool, req_pool_idx, seq_len=16)
|
||||
dense_slots = allocator.alloc(16)
|
||||
pool.req_to_token[req_pool_idx, :16] = dense_slots.to(torch.int32)
|
||||
compressed_cache = pool._aux_cache
|
||||
assert len(compressed_cache.free_slots) < len(compressed_cache.reserved_slots)
|
||||
|
||||
session = StreamingSession(
|
||||
SimpleNamespace(
|
||||
req_to_token_pool=pool,
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
page_size=1,
|
||||
)
|
||||
)
|
||||
session.slots["session-a"] = SessionSlot(
|
||||
req_pool_idx=req_pool_idx,
|
||||
kv=SimpleNamespace(kv_allocated_len=16),
|
||||
)
|
||||
|
||||
session.release_session("session-a")
|
||||
|
||||
assert req_pool_idx in pool.free_slots
|
||||
assert len(compressed_cache.free_slots) == len(compressed_cache.reserved_slots)
|
||||
|
||||
|
||||
def test_mamba_leak_diagnostic_does_not_report_reserved_slots():
|
||||
pool, _, _, allocator = make_pool_and_req(capacity=69)
|
||||
allocator.free_pages = torch.arange(6, 70, dtype=torch.int64)
|
||||
allocator.release_pages = torch.empty(0, dtype=torch.int64)
|
||||
pool.mamba_pool = SimpleNamespace(size=1)
|
||||
pool.mamba_allocator = SimpleNamespace(
|
||||
size=1,
|
||||
free_slots=torch.empty(0, dtype=torch.int64),
|
||||
)
|
||||
checker = SchedulerInvariantChecker(
|
||||
is_hybrid_swa=False,
|
||||
is_hybrid_ssm=True,
|
||||
disaggregation_mode=None,
|
||||
page_size=1,
|
||||
full_tokens_per_layer=None,
|
||||
swa_tokens_per_layer=None,
|
||||
max_total_num_tokens=64,
|
||||
tree_cache=SimpleNamespace(
|
||||
mamba_protected_size=lambda: 0,
|
||||
all_values_flatten=lambda: torch.empty(0, dtype=torch.int64),
|
||||
all_mamba_values_flatten=lambda: torch.empty(0, dtype=torch.int64),
|
||||
),
|
||||
token_to_kv_pool_allocator=allocator,
|
||||
req_to_token_pool=pool,
|
||||
pool_stats_observer=SimpleNamespace(
|
||||
session_held_mamba_slots=lambda: 0,
|
||||
),
|
||||
get_last_batch=lambda: None,
|
||||
get_running_batch=lambda: None,
|
||||
)
|
||||
|
||||
leak, message = checker._check_mamba_pool(
|
||||
SimpleNamespace(mamba_available_size=0, mamba_evictable_size=0)
|
||||
)
|
||||
|
||||
assert leak
|
||||
assert "leaked_full_pages" not in message
|
||||
assert "leaked_mamba_pages={1}" in message
|
||||
|
||||
|
||||
def test_partial_failure_rolls_back_and_free_releases_every_slot():
|
||||
"""A failed cache-level allocation must release slots allocated for other levels."""
|
||||
pool, req, req_pool_idx, allocator = make_pool_and_req(capacity=18)
|
||||
cache = pool._aux_cache
|
||||
|
||||
with pytest.raises(RuntimeError, match="out of reserved slots"):
|
||||
alloc_extend(pool, req_pool_idx, seq_len=16)
|
||||
|
||||
assert allocator.available_size() == 11
|
||||
assert len(cache.free_slots) == 7
|
||||
|
||||
allocator.capacity = 20
|
||||
allocator.clear()
|
||||
pool.reset_aux_cache_allocator()
|
||||
alloc_extend(pool, req_pool_idx, seq_len=16)
|
||||
assert allocator.available_size() == 12
|
||||
assert len(cache.free_slots) == 0
|
||||
|
||||
pool.free(req)
|
||||
assert req.req_pool_idx is None
|
||||
assert allocator.available_size() == 12
|
||||
assert len(cache.free_slots) == 8
|
||||
|
||||
|
||||
def test_allocator_reset_rebuilds_reserve():
|
||||
pool, _, _, allocator = make_pool_and_req()
|
||||
|
||||
allocator.clear()
|
||||
assert allocator.available_size() == 64
|
||||
|
||||
pool.reset_aux_cache_allocator()
|
||||
assert allocator.available_size() == 39
|
||||
assert len(pool._aux_cache.free_slots) == 25
|
||||
|
||||
|
||||
def test_attach_compressed_cache_is_idempotent():
|
||||
allocator = RecordingAllocator(capacity=64)
|
||||
pool = ReqToTokenPool(
|
||||
size=2,
|
||||
max_context_len=64,
|
||||
device="cpu",
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
attach_compressed_cache(
|
||||
pool,
|
||||
allocator,
|
||||
kernel_size=4,
|
||||
kernel_stride=2,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
cache = pool._aux_cache
|
||||
k1_table = pool.req_to_sparse_k1_token
|
||||
|
||||
attach_compressed_cache(
|
||||
pool,
|
||||
allocator,
|
||||
kernel_size=4,
|
||||
kernel_stride=2,
|
||||
enable_memory_saver=False,
|
||||
)
|
||||
|
||||
assert pool._aux_cache is cache
|
||||
assert pool.req_to_sparse_k1_token is k1_table
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
File diff suppressed because it is too large
Load Diff
@@ -18,6 +18,16 @@ class _FakeAllocator:
|
||||
self.freed.append(free_index.clone())
|
||||
|
||||
|
||||
class _FakeReqToTokenPool:
|
||||
def __init__(self, req_to_token):
|
||||
self.req_to_token = req_to_token
|
||||
self.free_slots = []
|
||||
|
||||
def free(self, req):
|
||||
self.free_slots.append(req.req_pool_idx)
|
||||
req.req_pool_idx = None
|
||||
|
||||
|
||||
class _FakeInnerCache:
|
||||
def __init__(self, req_to_token_pool, allocator, page_size, match_results=None):
|
||||
self.req_to_token_pool = req_to_token_pool
|
||||
@@ -85,7 +95,7 @@ def test_preabort_detaches_session_and_preserves_slot():
|
||||
"""Pre-aborted req (to_finish set before match_prefix) is detached from
|
||||
the session: session=None, abort_req() called. Slot stays intact."""
|
||||
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
||||
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
|
||||
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
||||
allocator = _FakeAllocator()
|
||||
inner = _FakeInnerCache(
|
||||
req_to_token_pool,
|
||||
@@ -133,7 +143,7 @@ def test_first_mid_abort_nukes_ephemeral_slot():
|
||||
slot is created from req state and nuked via release_session."""
|
||||
page_size = 1
|
||||
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
|
||||
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
|
||||
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
||||
allocator = _FakeAllocator()
|
||||
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
|
||||
tree_cache = StreamingSession(inner)
|
||||
@@ -159,7 +169,7 @@ def test_nth_mid_abort_nukes_session_slot():
|
||||
in req_nodes for next turn's re-prefill."""
|
||||
page_size = 1
|
||||
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
||||
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
|
||||
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
||||
allocator = _FakeAllocator()
|
||||
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size)
|
||||
tree_cache = StreamingSession(inner)
|
||||
@@ -197,7 +207,7 @@ def test_release_session_threads_mamba_skip_ids():
|
||||
from sglang.srt.mem_cache.unified_cache.components import ComponentType
|
||||
|
||||
req_to_token = torch.arange(256, dtype=torch.int32).reshape(2, 128)
|
||||
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
|
||||
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
||||
allocator = _FakeAllocator()
|
||||
inner = _FakeInnerCache(req_to_token_pool, allocator, page_size=1)
|
||||
tree_cache = StreamingSession(inner)
|
||||
@@ -235,7 +245,7 @@ def test_trim_overshoot_postcondition():
|
||||
"""
|
||||
page_size = 1
|
||||
req_to_token = torch.arange(128, dtype=torch.int32).reshape(1, 128)
|
||||
req_to_token_pool = SimpleNamespace(req_to_token=req_to_token, free_slots=[])
|
||||
req_to_token_pool = _FakeReqToTokenPool(req_to_token)
|
||||
allocator = _FakeAllocator()
|
||||
tree_cache = StreamingSession(
|
||||
_FakeInnerCache(req_to_token_pool, allocator, page_size)
|
||||
|
||||
+37
@@ -80,5 +80,42 @@ def test_split_full_attention_applies_model_wrapper_once():
|
||||
override.restore()
|
||||
|
||||
|
||||
def test_equal_resolved_backends_ignore_stale_global_backend():
|
||||
runner = SimpleNamespace(
|
||||
server_args=SimpleNamespace(
|
||||
attention_backend="global-test",
|
||||
speculative_attention_mode="prefill",
|
||||
),
|
||||
kv_cache_dtype=None,
|
||||
token_to_kv_pool=object(),
|
||||
req_to_token_pool=object(),
|
||||
init_new_workspace=None,
|
||||
)
|
||||
constructors = {
|
||||
"global-test": lambda _runner: _FakeBackend("global"),
|
||||
"resolved-test": lambda _runner: _FakeBackend("resolved"),
|
||||
}
|
||||
resolved = ResolvedAttentionBackendStr(
|
||||
decode="resolved-test",
|
||||
prefill="resolved-test",
|
||||
)
|
||||
|
||||
with (
|
||||
patch.dict(attention_backend_setup.ATTENTION_BACKENDS, constructors),
|
||||
patch.object(
|
||||
attention_backend_setup,
|
||||
"attn_backend_wrapper",
|
||||
side_effect=lambda _runner, backend: backend,
|
||||
),
|
||||
):
|
||||
result = attention_backend_setup._build_resolved_backend(
|
||||
model_runner=runner,
|
||||
resolved=resolved,
|
||||
init_new_workspace=False,
|
||||
)
|
||||
|
||||
assert result.name == "resolved"
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sys.exit(pytest.main([__file__, "-v"]))
|
||||
|
||||
@@ -22,6 +22,7 @@ from sglang.srt.arg_groups.overrides import (
|
||||
register_model_override,
|
||||
validate_declarations,
|
||||
)
|
||||
from sglang.srt.configs.minicpm import MiniCPMHybridConfig
|
||||
from sglang.srt.environ import envs
|
||||
from sglang.srt.runtime_context import (
|
||||
get_context,
|
||||
@@ -77,6 +78,7 @@ class TestModelOverridableWhitelist(CustomTestCase):
|
||||
"dcp_comm_backend",
|
||||
"dcp_replicate_q_proj",
|
||||
"disable_overlap_schedule",
|
||||
"disable_radix_cache",
|
||||
"uses_mamba_radix_cache",
|
||||
"mamba_radix_cache_strategy",
|
||||
"mamba_full_memory_ratio",
|
||||
@@ -294,6 +296,235 @@ class TestGoldenModelOverrides(_IsolatedPublish):
|
||||
"v_head_dim": 16,
|
||||
}
|
||||
|
||||
@staticmethod
|
||||
def _minicpm_overrides(
|
||||
architecture,
|
||||
*,
|
||||
sparse_attention=False,
|
||||
lightning_attention=False,
|
||||
attention_backend=None,
|
||||
prefill_attention_backend=None,
|
||||
decode_attention_backend=None,
|
||||
disaggregation_mode="null",
|
||||
enable_dp_attention=False,
|
||||
enable_hierarchical_cache=False,
|
||||
):
|
||||
args = SimpleNamespace(
|
||||
attention_backend=attention_backend,
|
||||
prefill_attention_backend=prefill_attention_backend,
|
||||
decode_attention_backend=decode_attention_backend,
|
||||
disaggregation_mode=disaggregation_mode,
|
||||
enable_dp_attention=enable_dp_attention,
|
||||
enable_hierarchical_cache=enable_hierarchical_cache,
|
||||
)
|
||||
args.is_attention_backend_not_set = lambda: all(
|
||||
backend is None
|
||||
for backend in (
|
||||
args.attention_backend,
|
||||
args.prefill_attention_backend,
|
||||
args.decode_attention_backend,
|
||||
)
|
||||
)
|
||||
mixer_types = []
|
||||
if sparse_attention:
|
||||
mixer_types.append("minicpm4")
|
||||
if lightning_attention:
|
||||
mixer_types.append("lightning-attn")
|
||||
if not mixer_types:
|
||||
mixer_types.append("minicpm4")
|
||||
declarations = collect_model_override_declarations(
|
||||
architecture,
|
||||
args,
|
||||
hf_config=MiniCPMHybridConfig(
|
||||
num_hidden_layers=len(mixer_types),
|
||||
num_attention_heads=1,
|
||||
num_key_value_heads=1,
|
||||
mixer_types=mixer_types,
|
||||
sparse_config={} if sparse_attention else None,
|
||||
),
|
||||
)
|
||||
return {
|
||||
field: value
|
||||
for _, declaration in declarations
|
||||
for field, value in declaration.items()
|
||||
}
|
||||
|
||||
def test_minicpm_disables_radix_cache_only_for_hybrid_layers(self):
|
||||
for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
|
||||
with self.subTest(architecture=architecture):
|
||||
self.assertNotIn(
|
||||
"disable_radix_cache",
|
||||
self._minicpm_overrides(architecture),
|
||||
)
|
||||
self.assertTrue(
|
||||
self._minicpm_overrides(architecture, sparse_attention=True)[
|
||||
"disable_radix_cache"
|
||||
]
|
||||
)
|
||||
self.assertTrue(
|
||||
self._minicpm_overrides(architecture, lightning_attention=True)[
|
||||
"disable_radix_cache"
|
||||
]
|
||||
)
|
||||
|
||||
def test_minicpm_rejects_dp_attention(self):
|
||||
for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
|
||||
with self.subTest(architecture=architecture):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"MiniCPM does not support DP attention",
|
||||
):
|
||||
self._minicpm_overrides(
|
||||
architecture,
|
||||
enable_dp_attention=True,
|
||||
)
|
||||
|
||||
def test_minicpm_rejects_hierarchical_cache_for_hybrid_models(self):
|
||||
for capability in ("sparse_attention", "lightning_attention"):
|
||||
with self.subTest(capability=capability):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"MiniCPM SALA does not support hierarchical cache",
|
||||
):
|
||||
self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
enable_hierarchical_cache=True,
|
||||
**{capability: True},
|
||||
)
|
||||
|
||||
def test_sparse_minicpm_defaults_to_sparse_attention_backend(self):
|
||||
with patch.object(
|
||||
overrides_module,
|
||||
"is_blackwell_supported",
|
||||
return_value=False,
|
||||
):
|
||||
for architecture in ("MiniCPMForCausalLM", "MiniCPMSALAForCausalLM"):
|
||||
with self.subTest(architecture=architecture):
|
||||
self.assertEqual(
|
||||
self._minicpm_overrides(
|
||||
architecture,
|
||||
sparse_attention=True,
|
||||
)["attention_backend"],
|
||||
"minicpm_flashattn",
|
||||
)
|
||||
|
||||
def test_minicpm_overrides_use_config_capabilities(self):
|
||||
args = SimpleNamespace(
|
||||
attention_backend=None,
|
||||
prefill_attention_backend=None,
|
||||
decode_attention_backend=None,
|
||||
disaggregation_mode="null",
|
||||
enable_dp_attention=False,
|
||||
enable_hierarchical_cache=False,
|
||||
is_attention_backend_not_set=lambda: True,
|
||||
)
|
||||
config = SimpleNamespace(
|
||||
has_minicpm_sparse_attention=True,
|
||||
has_lightning_layers=False,
|
||||
)
|
||||
|
||||
with patch.object(
|
||||
overrides_module, "is_blackwell_supported", return_value=False
|
||||
):
|
||||
overrides = overrides_module._minicpm_sala_overrides(args, config)
|
||||
|
||||
self.assertTrue(overrides["disable_radix_cache"])
|
||||
self.assertEqual(overrides["attention_backend"], "minicpm_flashattn")
|
||||
|
||||
def test_sparse_minicpm_defaults_to_flashinfer_on_blackwell(self):
|
||||
with patch.object(
|
||||
overrides_module,
|
||||
"is_blackwell_supported",
|
||||
return_value=True,
|
||||
):
|
||||
self.assertEqual(
|
||||
self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
)["attention_backend"],
|
||||
"minicpm_flashinfer",
|
||||
)
|
||||
|
||||
def test_minicpm_preserves_explicit_attention_backend(self):
|
||||
overrides = self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
attention_backend="fa3",
|
||||
)
|
||||
self.assertNotIn("attention_backend", overrides)
|
||||
|
||||
def test_sparse_minicpm_rejects_pd_disaggregation(self):
|
||||
for disaggregation_mode in ("prefill", "decode"):
|
||||
with self.subTest(disaggregation_mode=disaggregation_mode):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"MiniCPM sparse attention does not support PD disaggregation",
|
||||
):
|
||||
self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
disaggregation_mode=disaggregation_mode,
|
||||
)
|
||||
for backend_field in (
|
||||
"prefill_attention_backend",
|
||||
"decode_attention_backend",
|
||||
):
|
||||
with self.subTest(backend_field=backend_field):
|
||||
with self.assertRaisesRegex(
|
||||
ValueError,
|
||||
"MiniCPM sparse attention does not support PD disaggregation",
|
||||
):
|
||||
self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
disaggregation_mode="decode",
|
||||
**{backend_field: "minicpm_flashattn"},
|
||||
)
|
||||
|
||||
def test_minicpm_force_dense_uses_stock_attention_backend(self):
|
||||
with envs.SGLANG_MINICPM_FORCE_DENSE.override(True):
|
||||
self.assertNotIn(
|
||||
"attention_backend",
|
||||
self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
),
|
||||
)
|
||||
self.assertEqual(
|
||||
self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
attention_backend="minicpm_flashinfer",
|
||||
)["attention_backend"],
|
||||
"flashinfer",
|
||||
)
|
||||
with patch.object(
|
||||
overrides_module,
|
||||
"is_blackwell_supported",
|
||||
return_value=True,
|
||||
):
|
||||
self.assertEqual(
|
||||
self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
attention_backend="minicpm_flashattn",
|
||||
)["attention_backend"],
|
||||
"fa4",
|
||||
)
|
||||
with patch.object(
|
||||
overrides_module,
|
||||
"is_blackwell_supported",
|
||||
return_value=False,
|
||||
):
|
||||
split_overrides = self._minicpm_overrides(
|
||||
"MiniCPMSALAForCausalLM",
|
||||
sparse_attention=True,
|
||||
prefill_attention_backend="minicpm_flashattn",
|
||||
decode_attention_backend="minicpm_flashattn",
|
||||
)
|
||||
self.assertEqual(split_overrides["prefill_attention_backend"], "fa3")
|
||||
self.assertEqual(split_overrides["decode_attention_backend"], "fa3")
|
||||
|
||||
def _construct(self, arch, model_type, config_extra=None, **server_kwargs):
|
||||
from sglang.srt.server_args import ServerArgs
|
||||
|
||||
|
||||
Reference in New Issue
Block a user