[Spec Decode][Qwen4Exp] RecoverSSM for Qwen GDN and the PLE short conv
With --use-replayssm and speculative decoding, Qwen4Exp's Gated DeltaNet layers verify the draft window from a read-only checkpoint and keep a per-token replay record [correction(V), normalized key(K), gate(1)] in fp32. After sampling, one commit replays the accepted tokens and writes the recurrent state once (plus the block-boundary state in align mode), instead of snapshotting the state after every draft position. The PLE dilated short conv follows the same protocol (accepted-suffix compaction). num_speculative_blocks = 0, so no per-draft Mamba blocks are reserved. - GDN verify/commit kernels: model_executor/layers/mamba/gdn/recoverssm_gdn.py - metadata/backends: v1/attention/backends/gdn_recoverssm.py, ple_recoverssm.py - the verify reads strided q/k/v views of the conv output and writes in place - builders declare AttentionCGSupport.NEVER: FULL_AND_PIECEWISE falls back to PIECEWISE; the fallback now also counts breakable CUDA graphs (empty splitting_ops, with or without torch.compile) as piecewise attention instead of dropping to NONE - the generic Mamba metadata builder skips the Mamba2 ReplaySSM ring-buffer workspaces when --use-replayssm selects RecoverSSM (the PLE short conv uses that builder) - tests: GDN verify/commit vs the native kernels (fp32 and bf16 state, align boundaries, mixed and ragged batches, extreme gates), PLE vs token-by-token decode, config selection Signed-off-by: Jürgen Schmied <juergenschmied70@gmail.com> Co-authored-by: Claude <noreply@anthropic.com>
This commit is contained in:
Jürgen Schmied
@@ -0,0 +1,173 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""Configuration-level regression tests for Qwen4Exp GDN/PLE RecoverSSM selection.
|
||||
|
||||
``--use-replayssm`` with speculative decoding selects RecoverSSM for Qwen4Exp, as it
|
||||
does for Kimi-K3 KDA. The MTP drafter's derived config shares ``cache_config`` and must
|
||||
keep the choice. These tests run the real validator function against stand-in configs.
|
||||
"""
|
||||
|
||||
from types import SimpleNamespace
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
|
||||
from vllm.config.mamba import MambaBackendEnum
|
||||
from vllm.config.vllm import VllmConfig
|
||||
from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
|
||||
from vllm.v1.attention.backend import AttentionCGSupport
|
||||
|
||||
TARGET = "Qwen4ExpForConditionalGeneration"
|
||||
DRAFTER = "Qwen4ExpMTP"
|
||||
|
||||
|
||||
def _validator():
|
||||
validators = VllmConfig.__pydantic_decorators__.model_validators # type: ignore[attr-defined]
|
||||
return validators["validate_mamba_cached_kernel"].func
|
||||
|
||||
|
||||
def _config(
|
||||
arch: str,
|
||||
cache_config=None,
|
||||
num_spec: int = 3,
|
||||
use_replayssm: bool = True,
|
||||
supports_replayssm: bool = True,
|
||||
**overrides,
|
||||
):
|
||||
cfg = SimpleNamespace(
|
||||
num_speculative_tokens=num_spec,
|
||||
model_config=SimpleNamespace(
|
||||
architecture=arch, supports_replayssm=supports_replayssm
|
||||
),
|
||||
cache_config=cache_config
|
||||
or SimpleNamespace(
|
||||
use_recoverssm=False,
|
||||
use_replayssm=use_replayssm,
|
||||
mamba_cache_mode="align",
|
||||
),
|
||||
mamba_config=SimpleNamespace(
|
||||
enable_stochastic_rounding=False, backend=MambaBackendEnum.TRITON
|
||||
),
|
||||
use_v2_model_runner=True,
|
||||
parallel_config=SimpleNamespace(pipeline_parallel_size=1),
|
||||
kv_transfer_config=None,
|
||||
)
|
||||
for key, value in overrides.items():
|
||||
target, _, attr = key.partition("__")
|
||||
if attr:
|
||||
setattr(getattr(cfg, target), attr, value)
|
||||
else:
|
||||
setattr(cfg, target, value)
|
||||
return cfg
|
||||
|
||||
|
||||
def test_validator_is_registered():
|
||||
assert "validate_mamba_cached_kernel" in (
|
||||
VllmConfig.__pydantic_decorators__.model_validators # type: ignore[attr-defined]
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"arch", ["Qwen4ExpForConditionalGeneration", "Qwen4ExpForCausalLM"]
|
||||
)
|
||||
def test_replayssm_with_spec_selects_recoverssm(arch):
|
||||
cfg = _config(arch)
|
||||
_validator()(cfg)
|
||||
assert cfg.cache_config.use_recoverssm is True
|
||||
assert uses_recoverssm(cfg.cache_config, 3)
|
||||
|
||||
|
||||
def test_derived_drafter_keeps_the_choice_and_runs_the_checks():
|
||||
target = _config(TARGET)
|
||||
_validator()(target)
|
||||
drafter = _config(DRAFTER, cache_config=target.cache_config)
|
||||
_validator()(drafter)
|
||||
assert target.cache_config.use_recoverssm is True
|
||||
bad = _config(
|
||||
DRAFTER,
|
||||
cache_config=target.cache_config,
|
||||
parallel_config__pipeline_parallel_size=2,
|
||||
)
|
||||
with pytest.raises(ValueError, match="pipeline_parallel_size"):
|
||||
_validator()(bad)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("override", "value", "match"),
|
||||
[
|
||||
("parallel_config__pipeline_parallel_size", 2, "pipeline_parallel_size"),
|
||||
("mamba_config__enable_stochastic_rounding", True, "stochastic"),
|
||||
("mamba_config__backend", "flashinfer", "mamba-backend triton"),
|
||||
("use_v2_model_runner", False, "V2_MODEL_RUNNER"),
|
||||
("cache_config__mamba_cache_mode", "all", "none and align"),
|
||||
],
|
||||
)
|
||||
def test_runtime_checks_reject(override, value, match):
|
||||
cfg = _config(TARGET, **{override: value})
|
||||
with pytest.raises(ValueError, match=match):
|
||||
_validator()(cfg)
|
||||
|
||||
|
||||
def test_flag_unset_keeps_the_stock_path():
|
||||
cfg = _config(TARGET, use_replayssm=False)
|
||||
_validator()(cfg)
|
||||
assert cfg.cache_config.use_recoverssm is False
|
||||
assert not uses_recoverssm(cfg.cache_config, 3)
|
||||
|
||||
|
||||
def test_qwen4exp_replayssm_requires_speculative_decoding():
|
||||
with pytest.raises(ValueError, match="requires speculative decoding"):
|
||||
_validator()(_config(TARGET, num_spec=0))
|
||||
|
||||
|
||||
def test_architecture_without_replayssm_support_is_rejected():
|
||||
cfg = _config("Qwen3NextForCausalLM", supports_replayssm=False)
|
||||
with pytest.raises(ValueError, match="not supported for architecture"):
|
||||
_validator()(cfg)
|
||||
|
||||
|
||||
def test_uses_recoverssm_predicate():
|
||||
on = SimpleNamespace(use_recoverssm=True)
|
||||
assert uses_recoverssm(on, 3)
|
||||
assert not uses_recoverssm(on, 0)
|
||||
assert not uses_recoverssm(SimpleNamespace(use_recoverssm=False), 3)
|
||||
assert not uses_recoverssm(SimpleNamespace(), 3)
|
||||
|
||||
|
||||
def test_builders_fall_back_to_piecewise():
|
||||
from vllm.v1.attention.backends.gdn_recoverssm import GDNRecoverSSMMetadataBuilder
|
||||
from vllm.v1.attention.backends.ple_recoverssm import PleRecoverSSMMetadataBuilder
|
||||
|
||||
for builder in (GDNRecoverSSMMetadataBuilder, PleRecoverSSMMetadataBuilder):
|
||||
assert (
|
||||
builder.get_cudagraph_support(cast(Any, None), cast(Any, None))
|
||||
== AttentionCGSupport.NEVER
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("compile_mode", ["VLLM_COMPILE", "NONE"])
|
||||
@pytest.mark.parametrize(("breakable", "expected"), [("1", "PIECEWISE"), ("0", "NONE")])
|
||||
def test_never_support_falls_back_to_piecewise_with_breakable_cudagraphs(
|
||||
monkeypatch, compile_mode, breakable, expected
|
||||
):
|
||||
"""A backend without full-graph support must keep piecewise graphs when breakable
|
||||
CUDA graphs split at attention (splitting_ops is empty then, and torch.compile may
|
||||
be off)."""
|
||||
from vllm.config.compilation import (
|
||||
CompilationConfig,
|
||||
CompilationMode,
|
||||
CUDAGraphMode,
|
||||
)
|
||||
|
||||
monkeypatch.setenv("VLLM_USE_BREAKABLE_CUDAGRAPH", breakable)
|
||||
compilation_config = CompilationConfig(
|
||||
mode=CompilationMode[compile_mode],
|
||||
cudagraph_mode=CUDAGraphMode.FULL_AND_PIECEWISE,
|
||||
splitting_ops=[],
|
||||
)
|
||||
resolved = compilation_config.resolve_cudagraph_mode_and_sizes(
|
||||
AttentionCGSupport.NEVER,
|
||||
"GDNRecoverSSMAttentionBackend",
|
||||
uniform_decode_query_len=4,
|
||||
)
|
||||
assert resolved.name == expected
|
||||
@@ -0,0 +1,601 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""GDN RecoverSSM: verify against the native fused kernel, commit against the native
|
||||
per-token states, and the boundary checks (malformed metadata must raise ValueError
|
||||
before any kernel reads or writes out of bounds)."""
|
||||
|
||||
from typing import Any, cast
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import (
|
||||
GDNRecoverSSMCommitContext,
|
||||
gdn_recoverssm_verify,
|
||||
)
|
||||
from vllm.third_party.flash_linear_attention.ops.fused_sigmoid_gating import (
|
||||
fused_sigmoid_gating_delta_rule_update,
|
||||
)
|
||||
|
||||
pytestmark = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA")
|
||||
H, HV, K, V, T = 16, 48, 128, 128, 4
|
||||
QLENS = (4, 4, 3)
|
||||
|
||||
|
||||
def _inputs(nb: int = 10) -> dict:
|
||||
torch.manual_seed(0)
|
||||
dev = "cuda"
|
||||
tot = sum(QLENS)
|
||||
qsl = torch.tensor([0, 4, 8, 11], dtype=torch.int32, device=dev)
|
||||
return dict(
|
||||
A_log=(torch.rand(HV, device=dev) * 2 - 1).float(),
|
||||
a=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16),
|
||||
b=torch.randn(tot, HV, device=dev, dtype=torch.bfloat16),
|
||||
dt_bias=(torch.randn(HV, device=dev) * 0.5).bfloat16(),
|
||||
q=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16),
|
||||
k=torch.randn(1, tot, H, K, device=dev, dtype=torch.bfloat16),
|
||||
v=torch.randn(1, tot, HV, V, device=dev, dtype=torch.bfloat16),
|
||||
checkpoint_state=(torch.randn(nb, HV, V, K, device=dev) * 0.05).float(),
|
||||
replay_cache=torch.zeros(nb, HV, T, V + K + 1, device=dev),
|
||||
query_start_loc=qsl,
|
||||
state_indices=torch.tensor([2, 5, 7], dtype=torch.int32, device=dev),
|
||||
spec_query_len=T,
|
||||
)
|
||||
|
||||
|
||||
def _verify(inp: dict, **override):
|
||||
d = {**inp, **override}
|
||||
return gdn_recoverssm_verify(
|
||||
d["A_log"],
|
||||
d["a"],
|
||||
d["b"],
|
||||
d["dt_bias"],
|
||||
d["q"],
|
||||
d["k"],
|
||||
d["v"],
|
||||
checkpoint_state=d["checkpoint_state"],
|
||||
replay_cache=d["replay_cache"],
|
||||
query_start_loc=d["query_start_loc"],
|
||||
state_indices=d["state_indices"],
|
||||
spec_query_len=d["spec_query_len"],
|
||||
out=d.get("out"),
|
||||
)
|
||||
|
||||
|
||||
def _native(inp: dict):
|
||||
return fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=inp["A_log"],
|
||||
a=inp["a"],
|
||||
b=inp["b"],
|
||||
dt_bias=inp["dt_bias"],
|
||||
q=inp["q"],
|
||||
k=inp["k"],
|
||||
v=inp["v"],
|
||||
initial_state=inp["checkpoint_state"].clone(),
|
||||
inplace_final_state=False,
|
||||
cu_seqlens=inp["query_start_loc"].long(),
|
||||
ssm_state_indices=inp["state_indices"].long(),
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
)
|
||||
|
||||
|
||||
def _context(inp: dict, ckpt: torch.Tensor) -> GDNRecoverSSMCommitContext:
|
||||
nb = ckpt.shape[0]
|
||||
conv = [torch.zeros(nb, 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16)]
|
||||
return GDNRecoverSSMCommitContext.from_tensors(
|
||||
conv, [ckpt], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8
|
||||
)
|
||||
|
||||
|
||||
def test_verify_matches_native_and_keeps_the_checkpoint():
|
||||
inp = _inputs()
|
||||
before = inp["checkpoint_state"].clone()
|
||||
out = _verify(inp)
|
||||
ref, _ = _native(inp)
|
||||
torch.testing.assert_close(
|
||||
out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0
|
||||
)
|
||||
assert torch.equal(inp["checkpoint_state"], before)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("accepted", [1, 2, 3, 4])
|
||||
def test_commit_matches_native_per_token_state(accepted):
|
||||
inp = _inputs()
|
||||
_verify(inp)
|
||||
_, states = _native(inp)
|
||||
ckpt = inp["checkpoint_state"].clone()
|
||||
ctx = _context(inp, ckpt)
|
||||
acc = torch.tensor(
|
||||
[min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
ctx.commit(acc, inp["state_indices"], inp["query_start_loc"])
|
||||
for i, blk in enumerate(inp["state_indices"].tolist()):
|
||||
tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1
|
||||
torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6)
|
||||
others = [b for b in range(ckpt.shape[0]) if b not in inp["state_indices"].tolist()]
|
||||
assert torch.equal(ckpt[others], inp["checkpoint_state"][others])
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
"case", ["short_qsl", "capacity", "hv_lt_h", "a_log", "out", "replay", "device"]
|
||||
)
|
||||
def test_verify_rejects_malformed_inputs(case):
|
||||
inp = _inputs()
|
||||
over, match = {
|
||||
"short_qsl": (
|
||||
dict(query_start_loc=inp["query_start_loc"][:-1]),
|
||||
"query metadata",
|
||||
),
|
||||
"capacity": (
|
||||
dict(
|
||||
state_indices=inp["state_indices"][:2],
|
||||
query_start_loc=inp["query_start_loc"][:3],
|
||||
),
|
||||
"activation capacity",
|
||||
),
|
||||
"hv_lt_h": (
|
||||
dict(
|
||||
v=inp["v"][:, :, :8].contiguous(),
|
||||
a=inp["a"][:, :8].contiguous(),
|
||||
b=inp["b"][:, :8].contiguous(),
|
||||
),
|
||||
"positive multiple",
|
||||
),
|
||||
"a_log": (dict(A_log=torch.zeros(HV + 1, device="cuda")), "A_log"),
|
||||
"out": (
|
||||
dict(
|
||||
out=torch.empty(1, 11, HV, V + 1, device="cuda", dtype=torch.bfloat16)
|
||||
),
|
||||
"output shape",
|
||||
),
|
||||
"replay": (
|
||||
dict(replay_cache=torch.zeros(10, HV, T, V + K, device="cuda")),
|
||||
"replay buffer",
|
||||
),
|
||||
"device": (dict(A_log=inp["A_log"].cpu()), "same device"),
|
||||
}[case]
|
||||
with pytest.raises(ValueError, match=match):
|
||||
_verify(inp, **over)
|
||||
|
||||
|
||||
def test_commit_rejects_partial_or_short_metadata():
|
||||
inp = _inputs()
|
||||
ctx = _context(inp, inp["checkpoint_state"].clone())
|
||||
acc = torch.ones(3, dtype=torch.int32, device="cuda")
|
||||
bt = torch.zeros(3, 4, dtype=torch.int32, device="cuda")
|
||||
with pytest.raises(ValueError, match="align metadata is incomplete"):
|
||||
ctx.commit(
|
||||
acc,
|
||||
inp["state_indices"],
|
||||
inp["query_start_loc"],
|
||||
block_table=bt,
|
||||
mamba_block_size=16,
|
||||
)
|
||||
with pytest.raises(ValueError, match="commit metadata"):
|
||||
ctx.commit(acc, inp["state_indices"], inp["query_start_loc"][:-1])
|
||||
with pytest.raises(ValueError, match="request mapping"):
|
||||
ctx.commit(
|
||||
acc,
|
||||
inp["state_indices"],
|
||||
inp["query_start_loc"],
|
||||
request_indices=torch.zeros(2, dtype=torch.int32, device="cuda"),
|
||||
)
|
||||
|
||||
|
||||
def test_context_rejects_inconsistent_layers():
|
||||
inp = _inputs()
|
||||
ck = inp["checkpoint_state"]
|
||||
conv = torch.zeros(ck.shape[0], 8, 3 + T - 1, device="cuda", dtype=torch.bfloat16)
|
||||
with pytest.raises(ValueError, match="at least one layer"):
|
||||
GDNRecoverSSMCommitContext.from_tensors(
|
||||
[], [], [], spec_query_len=T, max_num_reqs=8
|
||||
)
|
||||
with pytest.raises(ValueError, match="differ"):
|
||||
GDNRecoverSSMCommitContext.from_tensors(
|
||||
[conv], [ck, ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8
|
||||
)
|
||||
with pytest.raises(ValueError, match="block count"):
|
||||
GDNRecoverSSMCommitContext.from_tensors(
|
||||
[conv[:5]], [ck], [inp["replay_cache"]], spec_query_len=T, max_num_reqs=8
|
||||
)
|
||||
|
||||
|
||||
# ------------------------------------------------------------------ review round 2:
|
||||
# mapping, align, extreme gates
|
||||
from vllm.v1.attention.backends.recoverssm_metadata import ( # noqa: E402
|
||||
recoverssm_request_indices,
|
||||
recoverssm_spec_rows,
|
||||
)
|
||||
|
||||
# (A_log range, a scale, b offset): drive g = -exp(A_log) * softplus(a + dt_bias)
|
||||
# and beta = sigmoid(b)
|
||||
GATES = {
|
||||
"typical": (0.0, 1.0, 0.0),
|
||||
"g_near_0": (-12.0, 1.0, 0.0),
|
||||
"g_very_negative": (2.5, 4.0, 0.0),
|
||||
"beta_near_0": (0.0, 1.0, -12.0),
|
||||
"beta_near_1": (0.0, 1.0, 12.0),
|
||||
}
|
||||
|
||||
|
||||
def _make(qlens, nb=16, seed=0, gate="typical"):
|
||||
g = torch.Generator(device="cuda").manual_seed(seed)
|
||||
rnd = lambda *s: torch.randn(*s, device="cuda", generator=g) # noqa: E731
|
||||
a_log_shift, ascale, boff = GATES[gate]
|
||||
tot = sum(qlens)
|
||||
qsl = torch.tensor(
|
||||
[0, *torch.tensor(qlens).cumsum(0).tolist()], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
return dict(
|
||||
A_log=(
|
||||
torch.rand(HV, device="cuda", generator=g) * 2 - 1 + a_log_shift
|
||||
).float(),
|
||||
a=(rnd(tot, HV) * ascale).bfloat16(),
|
||||
b=(rnd(tot, HV) + boff).bfloat16(),
|
||||
dt_bias=(rnd(HV) * 0.5).bfloat16(),
|
||||
q=rnd(1, tot, H, K).bfloat16(),
|
||||
k=rnd(1, tot, H, K).bfloat16(),
|
||||
v=rnd(1, tot, HV, V).bfloat16(),
|
||||
checkpoint_state=(rnd(nb, HV, V, K) * 0.05).float(),
|
||||
replay_cache=torch.zeros(nb, HV, T, V + K + 1, device="cuda"),
|
||||
query_start_loc=qsl,
|
||||
spec_query_len=T,
|
||||
)
|
||||
|
||||
|
||||
def test_spec_rows_mixed_batch():
|
||||
# rows: 0 spec K=3 | 1 prefill | 2 draft-less decode | 3 spec K=1 | 4 padded
|
||||
# zero-query | 5,6 beyond num_reqs
|
||||
is_prefilling = torch.tensor([False, True, False, False, False, False, False])
|
||||
qsl = torch.tensor([0, 4, 104, 105, 107, 107], dtype=torch.int32)
|
||||
drafts_in = torch.tensor(
|
||||
[3, -1, -1, 1, -1, 2, 0], dtype=torch.int32
|
||||
) # 5, 6: garbage padding
|
||||
drafts, rows = recoverssm_spec_rows(is_prefilling, qsl, drafts_in, num_reqs=5)
|
||||
assert rows.tolist() == [0, 2, 3]
|
||||
assert drafts.tolist() == [1, -1, 1, 1, -1, -1, -1]
|
||||
idx = recoverssm_request_indices(rows, 3, torch.device("cuda"))
|
||||
assert idx is not None and idx.tolist() == [0, 2, 3]
|
||||
assert (
|
||||
recoverssm_request_indices(torch.tensor([0, 1, 2]), 3, torch.device("cuda"))
|
||||
is None
|
||||
)
|
||||
with pytest.raises(ValueError, match="spec rows classified"):
|
||||
recoverssm_request_indices(rows, 2, torch.device("cuda"))
|
||||
|
||||
|
||||
def test_commit_maps_noncontiguous_requests_in_align_mode():
|
||||
"""Spec rows 0, 2, 3 of a 5-request batch; accepted counts, block table and
|
||||
num_computed are indexed by request row. Every spec request must commit its own
|
||||
state into its own block."""
|
||||
qlens = [4, 1, 2] # the spec rows' windows (K=3, draft-less, K=1)
|
||||
inp = _make(qlens, nb=40)
|
||||
bs = 8
|
||||
bt = torch.arange(1, 1 + 5 * 6, dtype=torch.int32, device="cuda").reshape(
|
||||
5, 6
|
||||
) # request row -> blocks
|
||||
nc = torch.tensor(
|
||||
[5, 0, 16, 9, 0], dtype=torch.int32, device="cuda"
|
||||
) # per request row
|
||||
req_rows = torch.tensor([0, 2, 3], dtype=torch.int32, device="cuda")
|
||||
src = torch.stack([bt[r, int(nc[r]) // bs] for r in req_rows.tolist()]).contiguous()
|
||||
inp["state_indices"] = src
|
||||
_verify(inp)
|
||||
_, states = _native(inp)
|
||||
ckpt = inp["checkpoint_state"].clone()
|
||||
ctx = _context(inp, ckpt)
|
||||
acc = torch.tensor(
|
||||
[3, 9, 1, 2, 9], dtype=torch.int32, device="cuda"
|
||||
) # rows 1, 4 must be ignored
|
||||
ctx.commit(
|
||||
acc,
|
||||
src,
|
||||
inp["query_start_loc"],
|
||||
request_indices=req_rows,
|
||||
block_table=bt,
|
||||
num_computed_tokens=nc,
|
||||
mamba_block_size=bs,
|
||||
)
|
||||
for i, r in enumerate(req_rows.tolist()):
|
||||
n = min(int(acc[r]), qlens[i])
|
||||
c0 = int(nc[r])
|
||||
final_blk = int(bt[r, (c0 + n) // bs])
|
||||
tok = int(inp["query_start_loc"][i]) + n - 1
|
||||
torch.testing.assert_close(ckpt[final_blk], states[tok], rtol=1e-5, atol=1e-6)
|
||||
|
||||
|
||||
# (num_computed, accepted, block size): boundary after 1 / after 2 / ending exactly on
|
||||
# it / crossing and continuing / accepted = 1 without crossing / the full window
|
||||
# crossing after 3
|
||||
ALIGN_CASES = [
|
||||
(3, 2, 4),
|
||||
(2, 3, 4),
|
||||
(0, 4, 4),
|
||||
(6, 4, 4),
|
||||
(1, 1, 4),
|
||||
(5, 4, 8),
|
||||
(7, 4, 8),
|
||||
(4, 4, 8),
|
||||
]
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES)
|
||||
def test_align_boundary_and_final_state_match_native(nc, accepted, bs):
|
||||
inp = _make([T], nb=12)
|
||||
bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6)
|
||||
src = bt[0, nc // bs].reshape(1).contiguous()
|
||||
inp["state_indices"] = src
|
||||
_verify(inp)
|
||||
_, states = _native(inp)
|
||||
ckpt = inp["checkpoint_state"].clone()
|
||||
ctx = _context(inp, ckpt)
|
||||
ctx.commit(
|
||||
torch.tensor([accepted], dtype=torch.int32, device="cuda"),
|
||||
src,
|
||||
inp["query_start_loc"],
|
||||
block_table=bt,
|
||||
num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"),
|
||||
mamba_block_size=bs,
|
||||
)
|
||||
next_boundary = (nc // bs + 1) * bs
|
||||
final_blk = int(bt[0, (nc + accepted) // bs])
|
||||
if (
|
||||
nc + accepted >= next_boundary
|
||||
): # crosses: the boundary state goes into the source block
|
||||
torch.testing.assert_close(
|
||||
ckpt[int(src)], states[next_boundary - nc - 1], rtol=1e-5, atol=1e-6
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
ckpt[final_blk], states[accepted - 1], rtol=1e-5, atol=1e-6
|
||||
)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("gate", list(GATES))
|
||||
@pytest.mark.parametrize("seed", [0, 1, 2])
|
||||
def test_commit_matches_native_under_extreme_gates(gate, seed):
|
||||
inp = _make(list(QLENS), nb=10, seed=seed, gate=gate)
|
||||
inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda")
|
||||
_verify(inp)
|
||||
_, states = _native(inp)
|
||||
for accepted in (1, 2, 3, 4):
|
||||
ckpt = inp["checkpoint_state"].clone()
|
||||
ctx = _context(inp, ckpt)
|
||||
acc = torch.tensor(
|
||||
[min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda"
|
||||
)
|
||||
ctx.commit(acc, inp["state_indices"], inp["query_start_loc"])
|
||||
for i, blk in enumerate(inp["state_indices"].tolist()):
|
||||
tok = int(inp["query_start_loc"][i]) + int(acc[i]) - 1
|
||||
torch.testing.assert_close(ckpt[blk], states[tok], rtol=1e-5, atol=1e-6)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("batch", [1, 2, 4])
|
||||
def test_verify_on_strided_views_writes_in_place(batch):
|
||||
"""The layer passes views of the packed conv output (token stride = the packed
|
||||
width) and the layer's output buffer; result and replay record must equal the
|
||||
contiguous-copy path."""
|
||||
torch.manual_seed(batch)
|
||||
H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6
|
||||
T = batch * SQ
|
||||
qd, vd = H * K, HV * V
|
||||
packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16)
|
||||
mixed = packed[:, : 2 * qd + vd]
|
||||
a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
|
||||
b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
|
||||
A_log = torch.randn(HV, device="cuda")
|
||||
dt_bias = torch.randn(HV, device="cuda")
|
||||
checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05
|
||||
qsl = torch.arange(0, T + 1, SQ, device="cuda", dtype=torch.int32)
|
||||
si = torch.arange(1, batch + 1, device="cuda", dtype=torch.int32)
|
||||
|
||||
def run(q, k, v, out=None):
|
||||
replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda")
|
||||
o = gdn_recoverssm_verify(
|
||||
A_log,
|
||||
a,
|
||||
b,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
checkpoint_state=checkpoint,
|
||||
replay_cache=replay,
|
||||
query_start_loc=qsl,
|
||||
state_indices=si,
|
||||
spec_query_len=SQ,
|
||||
out=out,
|
||||
)
|
||||
return o, replay
|
||||
|
||||
q, k, v = (
|
||||
x.contiguous().view(1, T, -1, d)
|
||||
for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V))
|
||||
)
|
||||
want_out, want_replay = run(q, k, v)
|
||||
|
||||
core = torch.zeros(T + 3, HV, V, device="cuda", dtype=torch.bfloat16)
|
||||
views = (
|
||||
mixed[:, :qd].view(1, T, -1, K),
|
||||
mixed[:, qd : 2 * qd].view(1, T, -1, K),
|
||||
mixed[:, 2 * qd :].view(1, T, -1, V),
|
||||
)
|
||||
got_out, got_replay = run(*views, out=core[:T].unsqueeze(0))
|
||||
assert got_out.data_ptr() == core.data_ptr()
|
||||
assert torch.equal(core[:T].unsqueeze(0), want_out)
|
||||
assert torch.equal(got_replay, want_replay)
|
||||
|
||||
|
||||
def test_verify_on_strided_views_with_variable_query_lengths():
|
||||
"""Token-strided views with ragged windows (query lengths 4, 1, 3), as MTP
|
||||
scheduling produces them."""
|
||||
torch.manual_seed(7)
|
||||
H, HV, K, V, SQ, NB = 16, 48, 128, 128, 4, 6
|
||||
qsl_list = [0, 4, 5, 8]
|
||||
T = qsl_list[-1]
|
||||
qd, vd = H * K, HV * V
|
||||
packed = torch.randn(T, 2 * qd + 2 * vd, device="cuda", dtype=torch.bfloat16)
|
||||
mixed = packed[:, : 2 * qd + vd]
|
||||
a = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
|
||||
b = torch.randn(T, HV, device="cuda", dtype=torch.bfloat16)
|
||||
A_log = torch.randn(HV, device="cuda")
|
||||
dt_bias = torch.randn(HV, device="cuda")
|
||||
checkpoint = torch.randn(NB, HV, V, K, device="cuda") * 0.05
|
||||
qsl = torch.tensor(qsl_list, device="cuda", dtype=torch.int32)
|
||||
si = torch.tensor([1, 3, 5], device="cuda", dtype=torch.int32)
|
||||
|
||||
def run(q, k, v, out=None):
|
||||
replay = torch.zeros(NB, HV, SQ, V + K + 1, device="cuda")
|
||||
o = gdn_recoverssm_verify(
|
||||
A_log,
|
||||
a,
|
||||
b,
|
||||
dt_bias,
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
checkpoint_state=checkpoint,
|
||||
replay_cache=replay,
|
||||
query_start_loc=qsl,
|
||||
state_indices=si,
|
||||
spec_query_len=SQ,
|
||||
out=out,
|
||||
)
|
||||
return o, replay
|
||||
|
||||
q, k, v = (
|
||||
x.contiguous().view(1, T, -1, d)
|
||||
for x, d in zip(torch.split(mixed, [qd, qd, vd], dim=-1), (K, K, V))
|
||||
)
|
||||
want_out, want_replay = run(q, k, v)
|
||||
core = torch.zeros(T + 2, HV, V, device="cuda", dtype=torch.bfloat16)
|
||||
views = (
|
||||
mixed[:, :qd].view(1, T, -1, K),
|
||||
mixed[:, qd : 2 * qd].view(1, T, -1, K),
|
||||
mixed[:, 2 * qd :].view(1, T, -1, V),
|
||||
)
|
||||
_, got_replay = run(*views, out=core[:T].unsqueeze(0))
|
||||
assert torch.equal(core[:T].unsqueeze(0), want_out)
|
||||
assert torch.equal(got_replay, want_replay)
|
||||
|
||||
|
||||
def test_layer_verify_qkv_views_and_fallback():
|
||||
"""The layer's guard: views for a [tokens, q|k|v] conv output with a contiguous
|
||||
last dimension, the contiguous-copy fallback for any other layout; same values
|
||||
either way."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from vllm.model_executor.layers.mamba.gdn.qwen_gdn_linear_attn import (
|
||||
QwenGatedDeltaNetAttention as Layer,
|
||||
)
|
||||
|
||||
H, HV, K, V, T = 4, 8, 16, 16, 5
|
||||
layer = SimpleNamespace(
|
||||
key_dim=H * K, value_dim=HV * V, tp_size=1, head_k_dim=K, head_v_dim=V
|
||||
)
|
||||
layer.rearrange_mixed_qkv = lambda x: Layer.rearrange_mixed_qkv(cast(Any, layer), x)
|
||||
width = 2 * H * K + HV * V
|
||||
packed = torch.randn(T, width + HV * V, dtype=torch.bfloat16) # q|k|v|z, like qkvz
|
||||
strided = packed[:, :width]
|
||||
q, k, v, is_view = Layer._recoverssm_verify_qkv(cast(Any, layer), strided)
|
||||
assert is_view
|
||||
assert q.data_ptr() == packed.data_ptr()
|
||||
rq, rk, rv = Layer.rearrange_mixed_qkv(cast(Any, layer), strided)
|
||||
assert torch.equal(q, rq) and torch.equal(k, rk) and torch.equal(v, rv)
|
||||
|
||||
# Column-major storage: last dimension not contiguous -> copies, same values.
|
||||
col_major = strided.t().contiguous().t()
|
||||
assert col_major.stride(-1) != 1
|
||||
q2, k2, v2, is_view2 = Layer._recoverssm_verify_qkv(cast(Any, layer), col_major)
|
||||
assert not is_view2
|
||||
assert torch.equal(q2, rq) and torch.equal(k2, rk) and torch.equal(v2, rv)
|
||||
|
||||
|
||||
# ------------------------------------------------------------- BF16 recurrent state
|
||||
# RecoverSSM upcasts the BF16 checkpoint, carries the accepted window in FP32 and rounds
|
||||
# once when it writes the committed state. The reference is the FP32 trajectory from the
|
||||
# same (already BF16-rounded) checkpoint; the committed state must equal its BF16
|
||||
# rounding up to one BF16 ulp, and be no further from it than the native path is.
|
||||
BF16_ULP = 2**-7
|
||||
|
||||
|
||||
def _bf16_case(qlens, nb=12, seed=0):
|
||||
inp = _make(qlens, nb=nb, seed=seed)
|
||||
inp["checkpoint_state"] = inp["checkpoint_state"].bfloat16()
|
||||
ref = dict(inp)
|
||||
ref["checkpoint_state"] = inp["checkpoint_state"].float()
|
||||
return inp, ref
|
||||
|
||||
|
||||
def test_verify_bf16_state_matches_native():
|
||||
inp, _ = _bf16_case(list(QLENS))
|
||||
inp["state_indices"] = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda")
|
||||
before = inp["checkpoint_state"].clone()
|
||||
out = _verify(inp)
|
||||
ref, _ = _native(inp)
|
||||
torch.testing.assert_close(
|
||||
out.float(), ref.reshape(out.shape).float(), rtol=0, atol=0
|
||||
)
|
||||
assert torch.equal(inp["checkpoint_state"], before)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("accepted", [1, 2, 3, 4])
|
||||
@pytest.mark.parametrize("seed", [0, 1])
|
||||
def test_commit_bf16_state_rounds_once(accepted, seed):
|
||||
inp, ref_inp = _bf16_case(list(QLENS), seed=seed)
|
||||
idx = torch.tensor([2, 5, 7], dtype=torch.int32, device="cuda")
|
||||
inp["state_indices"] = ref_inp["state_indices"] = idx
|
||||
_verify(inp)
|
||||
_, native_states = _native(inp)
|
||||
_, ref_states = _native(ref_inp)
|
||||
ckpt = inp["checkpoint_state"].clone()
|
||||
_context(inp, ckpt).commit(
|
||||
torch.tensor(
|
||||
[min(accepted, q) for q in QLENS], dtype=torch.int32, device="cuda"
|
||||
),
|
||||
idx,
|
||||
inp["query_start_loc"],
|
||||
)
|
||||
for i, blk in enumerate(idx.tolist()):
|
||||
tok = int(inp["query_start_loc"][i]) + min(accepted, QLENS[i]) - 1
|
||||
ref = ref_states[tok].float()
|
||||
got = ckpt[blk].float()
|
||||
torch.testing.assert_close(
|
||||
got, ref.bfloat16().float(), rtol=BF16_ULP, atol=1e-5
|
||||
)
|
||||
err_rssm = (got - ref).abs().mean()
|
||||
err_native = (native_states[tok].bfloat16().float() - ref).abs().mean()
|
||||
assert err_rssm <= err_native * 1.05 + 1e-7, (err_rssm, err_native)
|
||||
|
||||
|
||||
@pytest.mark.parametrize("nc,accepted,bs", ALIGN_CASES)
|
||||
def test_align_boundary_bf16_state(nc, accepted, bs):
|
||||
inp, ref_inp = _bf16_case([T])
|
||||
bt = torch.arange(2, 8, dtype=torch.int32, device="cuda").reshape(1, 6)
|
||||
src = bt[0, nc // bs].reshape(1).contiguous()
|
||||
inp["state_indices"] = ref_inp["state_indices"] = src
|
||||
_verify(inp)
|
||||
_, ref_states = _native(ref_inp)
|
||||
ckpt = inp["checkpoint_state"].clone()
|
||||
_context(inp, ckpt).commit(
|
||||
torch.tensor([accepted], dtype=torch.int32, device="cuda"),
|
||||
src,
|
||||
inp["query_start_loc"],
|
||||
block_table=bt,
|
||||
num_computed_tokens=torch.tensor([nc], dtype=torch.int32, device="cuda"),
|
||||
mamba_block_size=bs,
|
||||
)
|
||||
next_boundary = (nc // bs + 1) * bs
|
||||
final_blk = int(bt[0, (nc + accepted) // bs])
|
||||
if nc + accepted >= next_boundary:
|
||||
torch.testing.assert_close(
|
||||
ckpt[int(src)].float(),
|
||||
ref_states[next_boundary - nc - 1].float().bfloat16().float(),
|
||||
rtol=BF16_ULP,
|
||||
atol=1e-5,
|
||||
)
|
||||
torch.testing.assert_close(
|
||||
ckpt[final_blk].float(),
|
||||
ref_states[accepted - 1].float().bfloat16().float(),
|
||||
rtol=BF16_ULP,
|
||||
atol=1e-5,
|
||||
)
|
||||
@@ -0,0 +1,153 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""PLE short-conv RecoverSSM against the sequential reference.
|
||||
|
||||
Reference: the stock decode path fed the accepted tokens one at a time.
|
||||
RecoverSSM: one spec-mode conv over the whole window from the checkpoint
|
||||
(offset 0), then ``_PleConvCommit`` compacts the accepted suffix to the front.
|
||||
The conv history is pure data movement, so the states must be bit-identical,
|
||||
and the outputs of the accepted tokens must match as well.
|
||||
"""
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
|
||||
from vllm.models.qwen4_exp.nvidia.ops.ple import ple_conv
|
||||
from vllm.v1.attention.backends.ple_recoverssm import _PleConvCommit
|
||||
|
||||
C, KSIZE, DIL, SQ, NB = 512, 4, 3, 4, 8
|
||||
STATE_LEN = (KSIZE - 1) * DIL
|
||||
WIDTH = STATE_LEN + SQ - 1
|
||||
DEV = "cuda"
|
||||
|
||||
|
||||
def _state():
|
||||
if is_conv_state_dim_first():
|
||||
raw = torch.zeros(NB, C, WIDTH, dtype=torch.bfloat16, device=DEV)
|
||||
return raw, raw
|
||||
raw = torch.zeros(NB, WIDTH, C, dtype=torch.bfloat16, device=DEV)
|
||||
return raw, raw.transpose(-1, -2)
|
||||
|
||||
|
||||
def _i32(xs):
|
||||
return torch.tensor(xs, dtype=torch.int32, device=DEV)
|
||||
|
||||
|
||||
def _setup(seed, blocks):
|
||||
g = torch.Generator(device=DEV).manual_seed(seed)
|
||||
w = torch.randn(C, KSIZE, generator=g, device=DEV).to(torch.bfloat16)
|
||||
x = torch.randn(len(blocks) * SQ, C, generator=g, device=DEV).to(torch.bfloat16)
|
||||
outer = torch.randn_like(x)
|
||||
raw, view = _state()
|
||||
for b in blocks:
|
||||
view[b, :, :STATE_LEN] = torch.randn(C, STATE_LEN, generator=g, device=DEV).to(
|
||||
torch.bfloat16
|
||||
)
|
||||
return w, x, outer, raw, view
|
||||
|
||||
|
||||
def _sequential(w, x, outer, view, block, first, count):
|
||||
"""Stock decode, one token at a time, on a copy of the checkpoint block."""
|
||||
ref = view.clone()
|
||||
res = torch.zeros(count, C, dtype=torch.bfloat16, device=DEV)
|
||||
for i in range(count):
|
||||
t = first + i
|
||||
ple_conv(
|
||||
inputs=x[t : t + 1],
|
||||
residual=res[i : i + 1],
|
||||
conv_state=ref,
|
||||
conv_weights=w,
|
||||
state_indices=_i32([block]),
|
||||
outer_residual=outer[t : t + 1],
|
||||
mode="decode",
|
||||
dilation=DIL,
|
||||
)
|
||||
return ref[block, :, :STATE_LEN].clone(), res
|
||||
|
||||
|
||||
def _spec(w, x, outer, view, blocks):
|
||||
res = torch.zeros(len(blocks) * SQ, C, dtype=torch.bfloat16, device=DEV)
|
||||
ple_conv(
|
||||
inputs=x,
|
||||
residual=res,
|
||||
conv_state=view,
|
||||
conv_weights=w,
|
||||
state_indices=_i32(blocks),
|
||||
outer_residual=outer,
|
||||
mode="spec",
|
||||
dilation=DIL,
|
||||
query_start_loc=_i32(list(range(0, len(blocks) * SQ + 1, SQ))),
|
||||
num_accepted_tokens=_i32([1] * len(blocks)),
|
||||
spec_query_len=SQ,
|
||||
)
|
||||
return res
|
||||
|
||||
|
||||
@pytest.mark.parametrize("accepted", [1, 2, SQ])
|
||||
@pytest.mark.parametrize("seed", [0, 1])
|
||||
def test_ple_commit_matches_sequential(accepted, seed):
|
||||
block = 3
|
||||
w, x, outer, raw, view = _setup(seed, [block])
|
||||
want_state, want_out = _sequential(w, x, outer, view, block, 0, accepted)
|
||||
history = torch.cat([view[block, :, :STATE_LEN], x[:accepted].T], dim=1)
|
||||
|
||||
got_out = _spec(w, x, outer, view, [block])
|
||||
ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB)
|
||||
ctx.commit(_i32([accepted]), _i32([block]), _i32([0, SQ]))
|
||||
|
||||
got_state = view[block, :, :STATE_LEN]
|
||||
assert torch.equal(got_state, want_state)
|
||||
assert torch.equal(got_state, history[:, -STATE_LEN:])
|
||||
assert torch.equal(got_out[:accepted], want_out)
|
||||
|
||||
|
||||
@pytest.mark.parametrize(
|
||||
("num_computed", "accepted", "crosses"),
|
||||
[(6, 3, True), (6, 2, True), (1, 2, False), (7, SQ, True)],
|
||||
)
|
||||
def test_ple_commit_align_boundary(num_computed, accepted, crosses):
|
||||
block_size, src, nxt = 8, 2, 5
|
||||
w, x, outer, raw, view = _setup(7, [src])
|
||||
boundary_len = block_size - num_computed
|
||||
want_final, _ = _sequential(w, x, outer, view, src, 0, accepted)
|
||||
want_boundary, _ = _sequential(w, x, outer, view, src, 0, min(boundary_len, SQ))
|
||||
|
||||
_spec(w, x, outer, view, [src])
|
||||
ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB)
|
||||
ctx.commit(
|
||||
_i32([accepted]),
|
||||
_i32([src]),
|
||||
_i32([0, SQ]),
|
||||
block_table=_i32([[src, nxt]]),
|
||||
num_computed_tokens=_i32([num_computed]),
|
||||
mamba_block_size=block_size,
|
||||
)
|
||||
|
||||
final_block = nxt if num_computed + accepted >= block_size else src
|
||||
assert torch.equal(view[final_block, :, :STATE_LEN], want_final)
|
||||
if crosses:
|
||||
assert torch.equal(view[src, :, :STATE_LEN], want_boundary)
|
||||
|
||||
|
||||
def test_ple_commit_non_contiguous_request_rows():
|
||||
blocks, accepted = [1, 6], [3, 1]
|
||||
w, x, outer, raw, view = _setup(11, blocks)
|
||||
want = [
|
||||
_sequential(w, x, outer, view, b, i * SQ, a)[0]
|
||||
for i, (b, a) in enumerate(zip(blocks, accepted))
|
||||
]
|
||||
|
||||
_spec(w, x, outer, view, blocks)
|
||||
# Spec rows 0 and 1 belong to requests 1 and 3 of a four-request batch.
|
||||
counts = _i32([9, accepted[0], 9, accepted[1]])
|
||||
ctx = _PleConvCommit([raw], spec_query_len=SQ, max_num_reqs=NB)
|
||||
ctx.commit(
|
||||
counts,
|
||||
_i32(blocks),
|
||||
_i32([0, SQ, 2 * SQ]),
|
||||
request_indices=_i32([1, 3]),
|
||||
)
|
||||
|
||||
for b, ref in zip(blocks, want):
|
||||
assert torch.equal(view[b, :, :STATE_LEN], ref)
|
||||
@@ -1424,10 +1424,15 @@ class CompilationConfig:
|
||||
f"with {min_cg_attn_backend} backend (support: "
|
||||
f"{min_cg_support})"
|
||||
)
|
||||
if self.mode == CompilationMode.VLLM_COMPILE and (
|
||||
self.splitting_ops_contain_attention()
|
||||
or self.use_inductor_graph_partition
|
||||
):
|
||||
if (
|
||||
self.mode == CompilationMode.VLLM_COMPILE
|
||||
and (
|
||||
self.splitting_ops_contain_attention()
|
||||
or self.use_inductor_graph_partition
|
||||
)
|
||||
) or envs.VLLM_USE_BREAKABLE_CUDAGRAPH:
|
||||
# Breakable CUDA graphs break at the attention kernels at runtime,
|
||||
# with or without torch.compile.
|
||||
msg += (
|
||||
"; setting cudagraph_mode=PIECEWISE because "
|
||||
"attention is compiled piecewise"
|
||||
|
||||
+25
-4
@@ -70,6 +70,13 @@ else:
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
# Architectures whose GDN (and PLE) layers run RecoverSSM under --use-replayssm.
|
||||
_QWEN4EXP_RECOVERSSM_ARCHS = (
|
||||
"Qwen4ExpForCausalLM",
|
||||
"Qwen4ExpForConditionalGeneration",
|
||||
"Qwen4ExpMTP",
|
||||
)
|
||||
|
||||
# TODO(rocm): These models are either unsupported by MRV2 or slower with
|
||||
# MRV2 on AMD GPUs.
|
||||
ROCM_DEFAULT_MRV1_ARCHITECTURES = frozenset(
|
||||
@@ -3420,21 +3427,35 @@ class VllmConfig:
|
||||
@model_validator(mode="after")
|
||||
def validate_mamba_cached_kernel(self) -> "VllmConfig":
|
||||
if not self.cache_config.use_replayssm:
|
||||
self.cache_config.use_kda_recoverssm = False
|
||||
self.cache_config.use_recoverssm = False
|
||||
return self
|
||||
self.cache_config.use_kda_recoverssm = self.num_speculative_tokens > 0
|
||||
self.cache_config.use_recoverssm = self.num_speculative_tokens > 0
|
||||
|
||||
if self.model_config is not None and not self.model_config.supports_replayssm:
|
||||
raise ValueError(
|
||||
"--use-replayssm is not supported for architecture "
|
||||
f"{self.model_config.architecture!r}"
|
||||
)
|
||||
if self.cache_config.use_kda_recoverssm:
|
||||
if (
|
||||
self.model_config is not None
|
||||
and self.model_config.architecture in _QWEN4EXP_RECOVERSSM_ARCHS
|
||||
and not self.cache_config.use_recoverssm
|
||||
):
|
||||
# Qwen4Exp's GDN and PLE layers have no plain ReplaySSM path; with them the
|
||||
# flag selects RecoverSSM, which needs speculative decoding.
|
||||
raise ValueError(
|
||||
"--use-replayssm on Qwen4Exp selects GDN RecoverSSM and requires "
|
||||
"speculative decoding (num_speculative_tokens > 0)"
|
||||
)
|
||||
if self.cache_config.use_recoverssm:
|
||||
if self.model_config is not None and self.model_config.architecture not in (
|
||||
"KimiLinearForCausalLM",
|
||||
"KimiK3ForConditionalGeneration",
|
||||
*_QWEN4EXP_RECOVERSSM_ARCHS,
|
||||
):
|
||||
raise ValueError("RecoverSSM is only supported for Kimi-K3 KDA")
|
||||
raise ValueError(
|
||||
"RecoverSSM is only supported for Kimi-K3 KDA and Qwen4Exp GDN"
|
||||
)
|
||||
if self.mamba_config.enable_stochastic_rounding:
|
||||
raise ValueError(
|
||||
"RecoverSSM supports bfloat16/float32 "
|
||||
|
||||
@@ -370,10 +370,15 @@ class ChunkGatedDeltaRule(CustomOp):
|
||||
|
||||
@PluggableLayer.register("qwen_gated_delta_net_attention")
|
||||
class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
def get_state_shape(
|
||||
self,
|
||||
) -> tuple[tuple[int, ...], tuple[int, ...], tuple[int, ...], tuple[int, ...]]:
|
||||
return MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
def _uses_gdn_recoverssm(self) -> bool:
|
||||
"""One predicate for the replay state's shape, dtype and the attention
|
||||
backend."""
|
||||
from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
|
||||
|
||||
return uses_recoverssm(self.cache_config, self.num_spec)
|
||||
|
||||
def get_state_shape(self) -> tuple[tuple[int, ...], ...]:
|
||||
shapes = MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
self.tp_size,
|
||||
self.num_k_heads,
|
||||
self.num_v_heads,
|
||||
@@ -382,6 +387,32 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
self.conv_kernel_size,
|
||||
self.num_spec,
|
||||
)
|
||||
if self._uses_gdn_recoverssm():
|
||||
# per-token replay record [HV, spec_query_len, V + K + 1] fp32
|
||||
shapes = (
|
||||
*shapes,
|
||||
(
|
||||
self.num_v_heads // self.tp_size,
|
||||
self.num_spec + 1,
|
||||
self.head_v_dim + self.head_k_dim + 1,
|
||||
),
|
||||
)
|
||||
return shapes
|
||||
|
||||
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
|
||||
dtypes = super().get_state_dtype()
|
||||
if self._uses_gdn_recoverssm():
|
||||
dtypes = (*dtypes, torch.float32)
|
||||
return dtypes
|
||||
|
||||
def get_attn_backend(self):
|
||||
if self._uses_gdn_recoverssm():
|
||||
from vllm.v1.attention.backends.gdn_recoverssm import (
|
||||
GDNRecoverSSMAttentionBackend,
|
||||
)
|
||||
|
||||
return GDNRecoverSSMAttentionBackend
|
||||
return super().get_attn_backend()
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
@@ -530,6 +561,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
self.gdn_decode_kernel = "XPU"
|
||||
|
||||
self.enable_fused_gdn_decode = self.gdn_decode_kernel == "cuda"
|
||||
# RecoverSSM verify replaces the per-draft snapshot path (Triton decode
|
||||
# only)
|
||||
self.use_gdn_recoverssm = self._uses_gdn_recoverssm()
|
||||
if self.use_gdn_recoverssm:
|
||||
self.gdn_decode_kernel = "triton"
|
||||
self.enable_fused_gdn_decode = False
|
||||
self.enable_packed_recurrent_decode = False
|
||||
logger.info_once(
|
||||
"GDN RecoverSSM speculative verify active (spec_query_len %d)",
|
||||
self.num_spec + 1,
|
||||
)
|
||||
logger.info_once("GDN decode kernel: %s", self.gdn_decode_kernel)
|
||||
|
||||
compilation_config = get_current_vllm_config().compilation_config
|
||||
@@ -540,7 +582,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
def _fused_gdn_decode_unsupported_reason(
|
||||
self, vllm_config: VllmConfig
|
||||
) -> str | None:
|
||||
conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()
|
||||
conv_state_dtype, recurrent_state_dtype = self.get_state_dtype()[:2]
|
||||
if (
|
||||
self.gqa_interleaved_layout
|
||||
or self.head_k_dim != 128
|
||||
@@ -800,6 +842,25 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
|
||||
return mixed_qkv_out, z_out, b_out, a_out
|
||||
|
||||
def _recoverssm_verify_qkv(self, mixed_qkv):
|
||||
"""Q/K/V for the RecoverSSM verify.
|
||||
|
||||
The verify kernel reads token-strided heads, so for a 2-D [tokens, q|k|v] conv
|
||||
output with a contiguous last dimension it gets views instead of copies. Any
|
||||
other layout falls back to the contiguous copies of rearrange_mixed_qkv().
|
||||
Returns (query, key, value, is_view)."""
|
||||
if mixed_qkv.ndim != 2 or mixed_qkv.stride(-1) != 1:
|
||||
return (*self.rearrange_mixed_qkv(mixed_qkv), False)
|
||||
q_dim = self.key_dim // self.tp_size
|
||||
v_dim = self.value_dim // self.tp_size
|
||||
num_tokens = mixed_qkv.shape[0]
|
||||
query = mixed_qkv[:, :q_dim].view(1, num_tokens, -1, self.head_k_dim)
|
||||
key = mixed_qkv[:, q_dim : 2 * q_dim].view(1, num_tokens, -1, self.head_k_dim)
|
||||
value = mixed_qkv[:, 2 * q_dim : 2 * q_dim + v_dim].view(
|
||||
1, num_tokens, -1, self.head_v_dim
|
||||
)
|
||||
return query, key, value, True
|
||||
|
||||
def rearrange_mixed_qkv(self, mixed_qkv):
|
||||
"""Split packed qkv into contiguous (1, seq, heads, dim) tensors.
|
||||
|
||||
@@ -1085,7 +1146,7 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
dtype = qkv_or_qkvz.dtype
|
||||
num_k_heads = self.num_k_heads // self.tp_size
|
||||
num_v_heads = self.num_v_heads // self.tp_size
|
||||
_, state_dtype = self.get_state_dtype()
|
||||
state_dtype = self.get_state_dtype()[1] # 3 dtypes with the replay record
|
||||
|
||||
# All kernels use BT = chunk_size, so a single pass with T = chunk_size
|
||||
# is sufficient to populate every autotuner cache. Mirror the real
|
||||
@@ -1343,7 +1404,11 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
],
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
query_start_loc=spec_query_start_loc,
|
||||
max_query_len=spec_state_indices_tensor.size(-1),
|
||||
max_query_len=(
|
||||
self.num_spec + 1
|
||||
if getattr(self, "use_gdn_recoverssm", False)
|
||||
else spec_state_indices_tensor.size(-1)
|
||||
),
|
||||
validate_data=False,
|
||||
)
|
||||
|
||||
@@ -1380,7 +1445,17 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
else:
|
||||
mixed_qkv_non_spec = None
|
||||
|
||||
query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
|
||||
if (
|
||||
spec_sequence_masks is not None
|
||||
and mixed_qkv_spec is not None
|
||||
and self.use_gdn_recoverssm
|
||||
):
|
||||
query_spec, key_spec, value_spec, verify_views = (
|
||||
self._recoverssm_verify_qkv(mixed_qkv_spec)
|
||||
)
|
||||
else:
|
||||
verify_views = False
|
||||
query_spec, key_spec, value_spec = self.rearrange_mixed_qkv(mixed_qkv_spec)
|
||||
|
||||
# Split mixed non-spec-decode+prefill to process independently
|
||||
split_non_spec = (
|
||||
@@ -1443,7 +1518,43 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
# 2. Recurrent attention
|
||||
|
||||
# 2.1: Process the multi-query part
|
||||
if spec_sequence_masks is not None:
|
||||
if spec_sequence_masks is not None and getattr(
|
||||
self, "use_gdn_recoverssm", False
|
||||
):
|
||||
# verify the window off the checkpoint; the commit after sampling
|
||||
# writes the state once
|
||||
from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import (
|
||||
gdn_recoverssm_verify,
|
||||
)
|
||||
|
||||
_n = attn_metadata.num_spec_decodes
|
||||
assert spec_query_start_loc is not None
|
||||
assert spec_state_indices_tensor is not None
|
||||
core_attn_out_spec = gdn_recoverssm_verify(
|
||||
self.A_log,
|
||||
a_spec,
|
||||
b_spec,
|
||||
self.dt_bias,
|
||||
query_spec,
|
||||
key_spec,
|
||||
value_spec,
|
||||
checkpoint_state=ssm_state,
|
||||
replay_cache=self_kv_cache[2],
|
||||
query_start_loc=spec_query_start_loc[: _n + 1],
|
||||
state_indices=spec_state_indices_tensor[:_n, 0],
|
||||
spec_query_len=self.num_spec + 1,
|
||||
use_qk_l2norm_in_kernel=True,
|
||||
# spec-only batch: write straight into the layer output
|
||||
out=(
|
||||
core_attn_out[:num_actual_tokens].unsqueeze(0)
|
||||
if verify_views
|
||||
and mixed_qkv_non_spec is None
|
||||
and query_spec.shape[1] == num_actual_tokens
|
||||
else None
|
||||
),
|
||||
)
|
||||
last_recurrent_state = None
|
||||
elif spec_sequence_masks is not None:
|
||||
core_attn_out_spec, last_recurrent_state = (
|
||||
fused_sigmoid_gating_delta_rule_update(
|
||||
A_log=self.A_log,
|
||||
@@ -1558,7 +1669,8 @@ class QwenGatedDeltaNetAttention(GatedDeltaNetAttention):
|
||||
0, non_spec_token_indx, core_attn_out_non_spec.squeeze(0)
|
||||
)
|
||||
elif spec_sequence_masks is not None:
|
||||
core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
|
||||
if core_attn_out_spec.data_ptr() != core_attn_out.data_ptr():
|
||||
core_attn_out[:num_actual_tokens] = core_attn_out_spec.squeeze(0)
|
||||
else:
|
||||
core_attn_out[:num_actual_tokens] = core_attn_out_non_spec.squeeze(0)
|
||||
|
||||
|
||||
@@ -0,0 +1,695 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""RecoverSSM for Gated DeltaNet (Qwen-style GDN): speculative verify from one
|
||||
checkpoint without per-draft state snapshots, and accepted-state recovery after
|
||||
sampling.
|
||||
|
||||
GDN port of the Kimi-K3 KDA RecoverSSM path
|
||||
(vllm/models/kimi_k3/nvidia/ops/recoverssm.py). Differences from KDA: the gate is a
|
||||
scalar per value head (g = -exp(A_log) * softplus(a + dt_bias)), q/k heads are grouped
|
||||
over value heads (HV // H), and q/k are L2-normalised inside the kernel when requested.
|
||||
The verify step caches, per token, the delta-rule correction c_t = beta_t * (v_t - S'_t
|
||||
k_t) (V floats), the normalised key (K floats) and the log-decay g_t (1 float) in one
|
||||
fp32 "replay" record. Recovery folds the accepted prefix: S_n = exp(sum_t g_t) S_0 +
|
||||
sum_t exp(sum_{s>t} g_s) c_t k_t^T and writes the checkpoint once (plus the
|
||||
block-boundary state in align mode)."""
|
||||
|
||||
from collections.abc import Sequence
|
||||
from dataclasses import dataclass
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.model_executor.layers.mamba.ops.recoverssm import (
|
||||
compact_conv_state_kernel,
|
||||
prepare_commit_plan_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require
|
||||
from vllm.triton_utils import tl, triton
|
||||
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
|
||||
|
||||
def _require(cond: bool, msg: str) -> None:
|
||||
recoverssm_require(cond, msg, "GDN RecoverSSM")
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _gdn_recoverssm_verify_kernel(
|
||||
q_ptr,
|
||||
k_ptr,
|
||||
v_ptr,
|
||||
a_ptr,
|
||||
b_ptr,
|
||||
A_log_ptr,
|
||||
dt_bias_ptr,
|
||||
state_ptr,
|
||||
replay_ptr,
|
||||
out_ptr,
|
||||
query_start_loc_ptr,
|
||||
state_indices_ptr,
|
||||
scale,
|
||||
softplus_beta,
|
||||
softplus_threshold,
|
||||
null_block_id,
|
||||
stride_q_token,
|
||||
stride_k_token,
|
||||
stride_v_token,
|
||||
stride_a_token,
|
||||
stride_b_token,
|
||||
stride_state_block,
|
||||
stride_state_head,
|
||||
stride_state_v,
|
||||
stride_state_k,
|
||||
stride_replay_block,
|
||||
stride_replay_head,
|
||||
stride_replay_pos,
|
||||
stride_replay_dim,
|
||||
stride_out_token,
|
||||
stride_qsl,
|
||||
stride_si,
|
||||
H: tl.constexpr,
|
||||
HV: tl.constexpr,
|
||||
K: tl.constexpr,
|
||||
V: tl.constexpr,
|
||||
BK: tl.constexpr,
|
||||
BV: tl.constexpr,
|
||||
SPEC_QUERY_LEN: tl.constexpr,
|
||||
USE_QK_L2NORM: tl.constexpr,
|
||||
):
|
||||
pid_v = tl.program_id(0)
|
||||
pid_b = tl.program_id(1)
|
||||
pid_hv = tl.program_id(2)
|
||||
pid_h = pid_hv // (HV // H)
|
||||
bos = tl.load(query_start_loc_ptr + pid_b * stride_qsl).to(tl.int64)
|
||||
eos = tl.load(query_start_loc_ptr + (pid_b + 1) * stride_qsl).to(tl.int64)
|
||||
query_len = eos - bos
|
||||
state_idx = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64)
|
||||
offs_k = tl.arange(0, BK)
|
||||
offs_v = pid_v * BV + tl.arange(0, BV)
|
||||
mask_k = offs_k < K
|
||||
mask_v = offs_v < V
|
||||
mask_state = mask_v[:, None] & mask_k[None, :]
|
||||
if state_idx <= null_block_id:
|
||||
for t in tl.static_range(SPEC_QUERY_LEN):
|
||||
tl.store(
|
||||
out_ptr + (bos + t) * stride_out_token + pid_hv * V + offs_v,
|
||||
tl.zeros([BV], dtype=tl.float32),
|
||||
mask=(t < query_len) & mask_v,
|
||||
)
|
||||
return
|
||||
state = tl.load(
|
||||
state_ptr
|
||||
+ state_idx * stride_state_block
|
||||
+ pid_hv * stride_state_head
|
||||
+ offs_v[:, None] * stride_state_v
|
||||
+ offs_k[None, :] * stride_state_k,
|
||||
mask=mask_state,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
neg_A = -tl.exp(tl.load(A_log_ptr + pid_hv).to(tl.float32))
|
||||
dt_bias = tl.load(dt_bias_ptr + pid_hv).to(tl.float32)
|
||||
rec_base = (
|
||||
replay_ptr + state_idx * stride_replay_block + pid_hv * stride_replay_head
|
||||
)
|
||||
for t in tl.static_range(SPEC_QUERY_LEN):
|
||||
valid = t < query_len
|
||||
tok = bos + t
|
||||
q = tl.load(
|
||||
q_ptr + tok * stride_q_token + pid_h * K + offs_k,
|
||||
mask=valid & mask_k,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
k = tl.load(
|
||||
k_ptr + tok * stride_k_token + pid_h * K + offs_k,
|
||||
mask=valid & mask_k,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
v = tl.load(
|
||||
v_ptr + tok * stride_v_token + pid_hv * V + offs_v,
|
||||
mask=valid & mask_v,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
x = (
|
||||
tl.load(a_ptr + tok * stride_a_token + pid_hv, mask=valid, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
+ dt_bias
|
||||
)
|
||||
sp = tl.where(
|
||||
softplus_beta * x <= softplus_threshold,
|
||||
(1 / softplus_beta) * tl.log(1 + tl.exp(softplus_beta * x)),
|
||||
x,
|
||||
)
|
||||
g = neg_A * sp
|
||||
beta = tl.sigmoid(
|
||||
tl.load(b_ptr + tok * stride_b_token + pid_hv, mask=valid, other=0.0).to(
|
||||
tl.float32
|
||||
)
|
||||
)
|
||||
if USE_QK_L2NORM:
|
||||
q = q * tl.rsqrt(tl.sum(q * q) + 1e-6)
|
||||
k = k * tl.rsqrt(tl.sum(k * k) + 1e-6)
|
||||
q = q * scale
|
||||
new_state = state * tl.exp(g)
|
||||
corr = (v - tl.sum(new_state * k[None, :], 1)) * beta
|
||||
new_state = new_state + corr[:, None] * k[None, :]
|
||||
state = tl.where(valid, new_state, state)
|
||||
o = tl.sum(state * q[None, :], 1)
|
||||
tl.store(
|
||||
out_ptr + tok * stride_out_token + pid_hv * V + offs_v,
|
||||
o.to(out_ptr.dtype.element_ty),
|
||||
mask=valid & mask_v,
|
||||
)
|
||||
rec = rec_base + t * stride_replay_pos
|
||||
tl.store(rec + offs_v * stride_replay_dim, corr, mask=valid & mask_v)
|
||||
if pid_v == 0:
|
||||
tl.store(rec + (V + offs_k) * stride_replay_dim, k, mask=valid & mask_k)
|
||||
tl.store(rec + (V + K) * stride_replay_dim, g, mask=valid)
|
||||
|
||||
|
||||
def gdn_recoverssm_verify(
|
||||
A_log: torch.Tensor,
|
||||
a: torch.Tensor,
|
||||
b: torch.Tensor,
|
||||
dt_bias: torch.Tensor,
|
||||
q: torch.Tensor,
|
||||
k: torch.Tensor,
|
||||
v: torch.Tensor,
|
||||
*,
|
||||
checkpoint_state: torch.Tensor,
|
||||
replay_cache: torch.Tensor,
|
||||
query_start_loc: torch.Tensor,
|
||||
state_indices: torch.Tensor,
|
||||
spec_query_len: int,
|
||||
scale: float | None = None,
|
||||
softplus_beta: float = 1.0,
|
||||
softplus_threshold: float = 20.0,
|
||||
use_qk_l2norm_in_kernel: bool = True,
|
||||
out: torch.Tensor | None = None,
|
||||
) -> torch.Tensor:
|
||||
"""Verify a GDN speculative window against its checkpoint without modifying the
|
||||
checkpoint. q, k: [1, T, H, K]; v: [1, T, HV, V]; a, b: [T, HV] (or [1, T, HV]);
|
||||
checkpoint_state: [blocks, HV, V, K]; replay_cache: [blocks, HV, spec_query_len, V +
|
||||
K + 1] fp32. Returns out [1, T, HV, V]."""
|
||||
_require(
|
||||
q.ndim == 4 and q.shape[0] == 1, "q must have shape [1, tokens, heads, dim]"
|
||||
)
|
||||
_, total, H, K = q.shape
|
||||
_require(k.shape == q.shape, "q and k shapes differ")
|
||||
_require(
|
||||
v.ndim == 4 and tuple(v.shape[:2]) == (1, total),
|
||||
"v must have shape [1, tokens, value heads, dim]",
|
||||
)
|
||||
HV, V = v.shape[-2], v.shape[-1]
|
||||
_require(
|
||||
H > 0 and HV >= H and HV % H == 0,
|
||||
"value heads must be a positive multiple of the q/k heads",
|
||||
)
|
||||
_require(
|
||||
a.numel() == total * HV and b.numel() == total * HV,
|
||||
"gate or beta shape is incompatible",
|
||||
)
|
||||
a2 = a.reshape(total, HV)
|
||||
b2 = b.reshape(total, HV)
|
||||
_require(
|
||||
a2.stride(1) == 1 and b2.stride(1) == 1,
|
||||
"gate and beta heads must be contiguous",
|
||||
)
|
||||
_require(
|
||||
tuple(A_log.shape) == (HV,) and tuple(dt_bias.shape) == (HV,),
|
||||
"A_log or dt_bias shape is incompatible",
|
||||
)
|
||||
_require(
|
||||
all(t_.stride(-1) == 1 and t_.stride(-2) == K for t_ in (q, k)),
|
||||
"q and k heads must be contiguous",
|
||||
)
|
||||
_require(v.stride(-1) == 1 and v.stride(-2) == V, "v heads must be contiguous")
|
||||
_require(checkpoint_state.ndim == 4, "checkpoint must be four-dimensional")
|
||||
nb = checkpoint_state.shape[0]
|
||||
_require(
|
||||
tuple(checkpoint_state.shape[1:]) == (HV, V, K),
|
||||
"checkpoint shape is incompatible",
|
||||
)
|
||||
_rec = (nb, HV, spec_query_len, V + K + 1)
|
||||
_require(tuple(replay_cache.shape) == _rec, f"replay buffer needs shape {_rec}")
|
||||
_require(replay_cache.dtype == torch.float32, "replay buffer must use float32")
|
||||
_require(state_indices.ndim == 1, "state indices must be one-dimensional")
|
||||
_require(
|
||||
query_start_loc.ndim == 1
|
||||
and query_start_loc.shape[0] == state_indices.shape[0] + 1,
|
||||
"query metadata is incompatible",
|
||||
)
|
||||
_require(
|
||||
total <= state_indices.shape[0] * spec_query_len,
|
||||
"speculative decode input exceeds its activation capacity",
|
||||
)
|
||||
_dev = q.device
|
||||
_require(
|
||||
all(
|
||||
t_.device == _dev
|
||||
for t_ in (
|
||||
k,
|
||||
v,
|
||||
a,
|
||||
b,
|
||||
A_log,
|
||||
dt_bias,
|
||||
checkpoint_state,
|
||||
replay_cache,
|
||||
query_start_loc,
|
||||
state_indices,
|
||||
)
|
||||
if t_ is not None
|
||||
)
|
||||
and (out is None or out.device == _dev),
|
||||
"inputs must be on the same device",
|
||||
)
|
||||
if scale is None:
|
||||
scale = K**-0.5
|
||||
if out is None:
|
||||
out = torch.empty(1, total, HV, V, dtype=v.dtype, device=v.device)
|
||||
_require(
|
||||
tuple(out.shape) == (1, total, HV, V) and out.stride()[2:] == (V, 1),
|
||||
"output shape or layout is incompatible",
|
||||
)
|
||||
batch = state_indices.shape[0]
|
||||
if total == 0 or batch == 0:
|
||||
return out
|
||||
BK = triton.next_power_of_2(K)
|
||||
BV = min(triton.next_power_of_2(V), 32)
|
||||
grid = (triton.cdiv(V, BV), batch, HV)
|
||||
_gdn_recoverssm_verify_kernel[grid](
|
||||
q,
|
||||
k,
|
||||
v,
|
||||
a2,
|
||||
b2,
|
||||
A_log,
|
||||
dt_bias,
|
||||
checkpoint_state,
|
||||
replay_cache,
|
||||
out,
|
||||
query_start_loc,
|
||||
state_indices,
|
||||
scale,
|
||||
softplus_beta,
|
||||
softplus_threshold,
|
||||
NULL_BLOCK_ID,
|
||||
q.stride(1),
|
||||
k.stride(1),
|
||||
v.stride(1),
|
||||
a2.stride(0),
|
||||
b2.stride(0),
|
||||
checkpoint_state.stride(0),
|
||||
checkpoint_state.stride(1),
|
||||
checkpoint_state.stride(2),
|
||||
checkpoint_state.stride(3),
|
||||
replay_cache.stride(0),
|
||||
replay_cache.stride(1),
|
||||
replay_cache.stride(2),
|
||||
replay_cache.stride(3),
|
||||
out.stride(1),
|
||||
query_start_loc.stride(0),
|
||||
state_indices.stride(0),
|
||||
H=H,
|
||||
HV=HV,
|
||||
K=K,
|
||||
V=V,
|
||||
BK=BK,
|
||||
BV=BV,
|
||||
SPEC_QUERY_LEN=spec_query_len,
|
||||
USE_QK_L2NORM=use_qk_l2norm_in_kernel,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
return out
|
||||
|
||||
|
||||
@triton.jit
|
||||
def _commit_gdn_state_kernel(
|
||||
state_ref_ptr,
|
||||
state_base_addrs_ptr,
|
||||
state_block_strides_ptr,
|
||||
replay_ref_ptr,
|
||||
replay_base_addrs_ptr,
|
||||
replay_block_strides_ptr,
|
||||
state_indices_ptr,
|
||||
commit_lens_ptr,
|
||||
final_state_indices_ptr,
|
||||
boundary_state_indices_ptr,
|
||||
boundary_recovery_lens_ptr,
|
||||
null_block_id,
|
||||
stride_state_head,
|
||||
stride_state_v,
|
||||
stride_state_k,
|
||||
stride_replay_head,
|
||||
stride_replay_pos,
|
||||
stride_replay_dim,
|
||||
stride_si,
|
||||
K: tl.constexpr,
|
||||
V: tl.constexpr,
|
||||
BK: tl.constexpr,
|
||||
BV: tl.constexpr,
|
||||
NUM_HEADS: tl.constexpr,
|
||||
ALIGN_MODE: tl.constexpr,
|
||||
):
|
||||
pid_v = tl.program_id(0)
|
||||
pid_b = tl.program_id(1)
|
||||
pid_lh = tl.program_id(2)
|
||||
pid_l = pid_lh // NUM_HEADS
|
||||
pid_h = pid_lh % NUM_HEADS
|
||||
src = tl.load(state_indices_ptr + pid_b * stride_si).to(tl.int64)
|
||||
if src <= null_block_id:
|
||||
return
|
||||
n = tl.load(commit_lens_ptr + pid_b)
|
||||
if n == 0:
|
||||
return
|
||||
final_idx = tl.load(final_state_indices_ptr + pid_b).to(tl.int64)
|
||||
boundary_idx = tl.load(boundary_state_indices_ptr + pid_b).to(tl.int64)
|
||||
boundary_len = tl.load(boundary_recovery_lens_ptr + pid_b)
|
||||
if final_idx <= null_block_id:
|
||||
return
|
||||
state_ptr = tl.load(state_base_addrs_ptr + pid_l).to(
|
||||
tl.pointer_type(state_ref_ptr.dtype.element_ty)
|
||||
)
|
||||
sbs = tl.load(state_block_strides_ptr + pid_l)
|
||||
replay_ptr = tl.load(replay_base_addrs_ptr + pid_l).to(
|
||||
tl.pointer_type(replay_ref_ptr.dtype.element_ty)
|
||||
)
|
||||
rbs = tl.load(replay_block_strides_ptr + pid_l)
|
||||
rec_base = replay_ptr + src * rbs + pid_h * stride_replay_head
|
||||
offs_k = tl.arange(0, BK)
|
||||
offs_v = pid_v * BV + tl.arange(0, BV)
|
||||
mask_k = offs_k < K
|
||||
mask_v = offs_v < V
|
||||
mask_state = mask_v[:, None] & mask_k[None, :]
|
||||
sp = offs_v[:, None] * stride_state_v + offs_k[None, :] * stride_state_k
|
||||
s0 = tl.load(
|
||||
state_ptr + src * sbs + pid_h * stride_state_head + sp,
|
||||
mask=mask_state,
|
||||
other=0.0,
|
||||
).to(tl.float32)
|
||||
fin_decay = 1.0
|
||||
fin_corr = tl.zeros([BV, BK], tl.float32)
|
||||
bnd_decay = 1.0
|
||||
bnd_corr = tl.zeros([BV, BK], tl.float32)
|
||||
for r in range(n):
|
||||
t = n - r - 1
|
||||
rec = rec_base + t * stride_replay_pos
|
||||
c = tl.load(rec + offs_v * stride_replay_dim, mask=mask_v, other=0.0)
|
||||
kk = tl.load(rec + (V + offs_k) * stride_replay_dim, mask=mask_k, other=0.0)
|
||||
g = tl.load(rec + (V + K) * stride_replay_dim)
|
||||
upd = c[:, None] * kk[None, :]
|
||||
dec = tl.exp(g)
|
||||
fin_corr += upd * fin_decay
|
||||
fin_decay *= dec
|
||||
if ALIGN_MODE:
|
||||
before = t < boundary_len
|
||||
bnd_corr += tl.where(before, upd * bnd_decay, 0.0)
|
||||
bnd_decay *= tl.where(before, dec, 1.0)
|
||||
if ALIGN_MODE:
|
||||
bs = s0 * bnd_decay + bnd_corr
|
||||
tl.store(
|
||||
state_ptr + boundary_idx * sbs + pid_h * stride_state_head + sp,
|
||||
bs.to(state_ref_ptr.dtype.element_ty),
|
||||
mask=mask_state & (boundary_idx > null_block_id),
|
||||
)
|
||||
fs = s0 * fin_decay + fin_corr
|
||||
tl.store(
|
||||
state_ptr + final_idx * sbs + pid_h * stride_state_head + sp,
|
||||
fs.to(state_ref_ptr.dtype.element_ty),
|
||||
mask=mask_state,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GDNRecoverSSMCommitContext:
|
||||
conv_states: tuple[torch.Tensor, ...]
|
||||
conv_state_base_addrs: torch.Tensor
|
||||
conv_state_block_strides: torch.Tensor
|
||||
conv_state_dim_strides: torch.Tensor
|
||||
conv_state_token_strides: torch.Tensor
|
||||
conv_history_len: int
|
||||
checkpoints: tuple[torch.Tensor, ...]
|
||||
state_base_addrs: torch.Tensor
|
||||
state_block_strides: torch.Tensor
|
||||
replays: tuple[torch.Tensor, ...]
|
||||
replay_base_addrs: torch.Tensor
|
||||
replay_block_strides: torch.Tensor
|
||||
commit_lens: torch.Tensor
|
||||
final_state_indices: torch.Tensor
|
||||
boundary_state_indices: torch.Tensor
|
||||
boundary_recovery_lens: torch.Tensor
|
||||
spec_query_len: int
|
||||
|
||||
@classmethod
|
||||
def from_tensors(
|
||||
cls,
|
||||
conv_states: Sequence[torch.Tensor],
|
||||
checkpoints: Sequence[torch.Tensor],
|
||||
replays: Sequence[torch.Tensor],
|
||||
*,
|
||||
spec_query_len: int,
|
||||
max_num_reqs: int,
|
||||
conv_dim_first: bool = True,
|
||||
) -> "GDNRecoverSSMCommitContext":
|
||||
_require(len(checkpoints) > 0, "commit requires at least one layer")
|
||||
_require(
|
||||
len(conv_states) == len(checkpoints) == len(replays),
|
||||
"conv, state and replay lists differ",
|
||||
)
|
||||
if not conv_dim_first:
|
||||
conv_states = [s.transpose(-1, -2) for s in conv_states]
|
||||
ref = checkpoints[0]
|
||||
_require(ref.ndim == 4, "checkpoint must be four-dimensional")
|
||||
nb, HV, V, K = ref.shape
|
||||
_dev = ref.device
|
||||
for s in checkpoints:
|
||||
_require(
|
||||
s.shape == ref.shape
|
||||
and s.dtype == ref.dtype
|
||||
and s.stride()[1:] == ref.stride()[1:]
|
||||
and s.device == _dev,
|
||||
"layers need matching checkpoints",
|
||||
)
|
||||
_rec = (nb, HV, spec_query_len, V + K + 1)
|
||||
for r in replays:
|
||||
_require(
|
||||
tuple(r.shape) == _rec
|
||||
and r.dtype == torch.float32
|
||||
and r.stride()[1:] == replays[0].stride()[1:]
|
||||
and r.device == _dev,
|
||||
f"layers need matching float32 replay buffers of shape {_rec}",
|
||||
)
|
||||
conv_ref = conv_states[0]
|
||||
_require(
|
||||
conv_ref.ndim == 3 and conv_ref.shape[0] == nb,
|
||||
"conv state must be [blocks, dim, window] with the checkpoint's "
|
||||
"block count",
|
||||
)
|
||||
for c in conv_states:
|
||||
_require(
|
||||
c.shape == conv_ref.shape
|
||||
and c.dtype == conv_ref.dtype
|
||||
and c.stride()[1:] == conv_ref.stride()[1:]
|
||||
and c.device == _dev,
|
||||
"layers need matching conv states",
|
||||
)
|
||||
conv_dim, conv_len = conv_ref.shape[1:]
|
||||
hist = conv_len - spec_query_len + 1
|
||||
_require(hist > 0, "conv state is shorter than its window")
|
||||
dev = ref.device
|
||||
addr = lambda ts: torch.tensor(
|
||||
[t.data_ptr() for t in ts], dtype=torch.int64, device=dev
|
||||
)
|
||||
bstr = lambda ts: torch.tensor(
|
||||
[t.stride(0) for t in ts], dtype=torch.int64, device=dev
|
||||
)
|
||||
z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev)
|
||||
return cls(
|
||||
tuple(conv_states),
|
||||
addr(conv_states),
|
||||
bstr(conv_states),
|
||||
torch.tensor(
|
||||
[s.stride(1) for s in conv_states], dtype=torch.int64, device=dev
|
||||
),
|
||||
torch.tensor(
|
||||
[s.stride(2) for s in conv_states], dtype=torch.int64, device=dev
|
||||
),
|
||||
hist,
|
||||
tuple(checkpoints),
|
||||
addr(checkpoints),
|
||||
bstr(checkpoints),
|
||||
tuple(replays),
|
||||
addr(replays),
|
||||
bstr(replays),
|
||||
z(),
|
||||
z(),
|
||||
z(),
|
||||
z(),
|
||||
spec_query_len,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def create(
|
||||
cls, layers: Sequence[Any], *, spec_query_len: int, max_num_reqs: int
|
||||
) -> "GDNRecoverSSMCommitContext":
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
|
||||
|
||||
if any(len(layer.kv_cache) != 3 for layer in layers):
|
||||
raise ValueError("GDN RecoverSSM pages must contain conv, state and replay")
|
||||
return cls.from_tensors(
|
||||
[layer.kv_cache[0] for layer in layers],
|
||||
[layer.kv_cache[1] for layer in layers],
|
||||
[layer.kv_cache[2] for layer in layers],
|
||||
spec_query_len=spec_query_len,
|
||||
max_num_reqs=max_num_reqs,
|
||||
conv_dim_first=is_conv_state_dim_first(),
|
||||
)
|
||||
|
||||
def commit(
|
||||
self,
|
||||
num_accepted_tokens: torch.Tensor,
|
||||
state_indices: torch.Tensor,
|
||||
query_start_loc: torch.Tensor,
|
||||
request_indices: torch.Tensor | None = None,
|
||||
block_table: torch.Tensor | None = None,
|
||||
num_computed_tokens: torch.Tensor | None = None,
|
||||
mamba_block_size: int | None = None,
|
||||
commit_conv: bool = True,
|
||||
) -> None:
|
||||
"""Fold accepted GDN and convolution inputs into every layer's checkpoint."""
|
||||
batch = state_indices.shape[0]
|
||||
if batch == 0:
|
||||
return
|
||||
_require(state_indices.ndim == 1, "state indices must be one-dimensional")
|
||||
_require(
|
||||
batch <= self.commit_lens.shape[0], "commit batch exceeds its plan capacity"
|
||||
)
|
||||
_require(
|
||||
query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1,
|
||||
"commit metadata is incompatible",
|
||||
)
|
||||
_require(
|
||||
request_indices is None or request_indices.shape[0] >= batch,
|
||||
"request mapping is too short",
|
||||
)
|
||||
_require(
|
||||
num_accepted_tokens.ndim == 1
|
||||
and (request_indices is not None or num_accepted_tokens.shape[0] >= batch),
|
||||
"accepted-token counts are too short",
|
||||
)
|
||||
_align = (block_table, num_computed_tokens, mamba_block_size)
|
||||
_require(
|
||||
all(x is None for x in _align) or all(x is not None for x in _align),
|
||||
"align metadata is incomplete",
|
||||
)
|
||||
_require(
|
||||
mamba_block_size is None or mamba_block_size >= self.spec_query_len,
|
||||
"align block size must cover one speculative window",
|
||||
)
|
||||
_require(
|
||||
block_table is None or block_table.ndim == 2,
|
||||
"block table must be two-dimensional",
|
||||
)
|
||||
_dev = self.checkpoints[0].device
|
||||
_require(
|
||||
all(
|
||||
t_ is None or t_.device == _dev
|
||||
for t_ in (
|
||||
num_accepted_tokens,
|
||||
state_indices,
|
||||
query_start_loc,
|
||||
request_indices,
|
||||
block_table,
|
||||
num_computed_tokens,
|
||||
)
|
||||
),
|
||||
"commit inputs must be on the same device",
|
||||
)
|
||||
bt_stride = (0, 0) if block_table is None else block_table.stride()
|
||||
prepare_commit_plan_kernel[(batch,)](
|
||||
num_accepted_tokens,
|
||||
request_indices,
|
||||
state_indices,
|
||||
query_start_loc,
|
||||
block_table,
|
||||
num_computed_tokens,
|
||||
self.commit_lens,
|
||||
self.final_state_indices,
|
||||
self.boundary_state_indices,
|
||||
self.boundary_recovery_lens,
|
||||
NULL_BLOCK_ID,
|
||||
mamba_block_size or 1,
|
||||
block_table.shape[1] if block_table is not None else 1,
|
||||
num_accepted_tokens.stride(0),
|
||||
request_indices.stride(0) if request_indices is not None else 0,
|
||||
state_indices.stride(0),
|
||||
query_start_loc.stride(0),
|
||||
bt_stride[0],
|
||||
bt_stride[1],
|
||||
0 if num_computed_tokens is None else num_computed_tokens.stride(0),
|
||||
SPEC_QUERY_LEN=self.spec_query_len,
|
||||
num_warps=1,
|
||||
)
|
||||
num_layers = len(self.checkpoints)
|
||||
if commit_conv:
|
||||
conv_ref = self.conv_states[0]
|
||||
conv_dim = conv_ref.shape[1]
|
||||
compact_conv_state_kernel[(triton.cdiv(conv_dim, 256), batch, num_layers)](
|
||||
conv_ref,
|
||||
self.conv_state_base_addrs,
|
||||
self.conv_state_block_strides,
|
||||
self.conv_state_dim_strides,
|
||||
self.conv_state_token_strides,
|
||||
state_indices,
|
||||
self.commit_lens,
|
||||
self.final_state_indices,
|
||||
self.boundary_state_indices,
|
||||
self.boundary_recovery_lens,
|
||||
NULL_BLOCK_ID,
|
||||
conv_dim,
|
||||
self.conv_history_len,
|
||||
state_indices.stride(0),
|
||||
BLOCK_D=256,
|
||||
BLOCK_HISTORY=triton.next_power_of_2(self.conv_history_len),
|
||||
ALIGN_MODE=block_table is not None,
|
||||
num_warps=4,
|
||||
)
|
||||
ref = self.checkpoints[0]
|
||||
_, HV, V, K = ref.shape
|
||||
BK = triton.next_power_of_2(K)
|
||||
BV = min(triton.next_power_of_2(V), 32)
|
||||
rr = self.replays[0]
|
||||
_commit_gdn_state_kernel[(triton.cdiv(V, BV), batch, num_layers * HV)](
|
||||
ref,
|
||||
self.state_base_addrs,
|
||||
self.state_block_strides,
|
||||
rr,
|
||||
self.replay_base_addrs,
|
||||
self.replay_block_strides,
|
||||
state_indices,
|
||||
self.commit_lens,
|
||||
self.final_state_indices,
|
||||
self.boundary_state_indices,
|
||||
self.boundary_recovery_lens,
|
||||
NULL_BLOCK_ID,
|
||||
ref.stride(1),
|
||||
ref.stride(2),
|
||||
ref.stride(3),
|
||||
rr.stride(1),
|
||||
rr.stride(2),
|
||||
rr.stride(3),
|
||||
state_indices.stride(0),
|
||||
K=K,
|
||||
V=V,
|
||||
BK=BK,
|
||||
BV=BV,
|
||||
NUM_HEADS=HV,
|
||||
ALIGN_MODE=block_table is not None,
|
||||
num_warps=4,
|
||||
num_stages=2,
|
||||
)
|
||||
|
||||
|
||||
__all__ = ["GDNRecoverSSMCommitContext", "gdn_recoverssm_verify"]
|
||||
@@ -85,6 +85,7 @@ MambaStateShapes: TypeAlias = (
|
||||
| tuple[tuple[int, int, int]]
|
||||
| tuple[tuple[int, int], tuple[int, int]]
|
||||
| tuple[tuple[int, int], tuple[int, int, int]]
|
||||
| tuple[tuple[int, int], tuple[int, int, int], tuple[int, int, int]]
|
||||
| tuple[
|
||||
tuple[int, int],
|
||||
tuple[int, int, int],
|
||||
|
||||
@@ -32,11 +32,13 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
from vllm.model_executor.models.interfaces import (
|
||||
HasInnerState,
|
||||
IsHybrid,
|
||||
MambaStateShapes,
|
||||
MixtureOfExperts,
|
||||
MultiModalEmbeddings,
|
||||
SupportsLoRA,
|
||||
SupportsMRoPE,
|
||||
SupportsPP,
|
||||
SupportsReplaySSM,
|
||||
_require_is_multimodal,
|
||||
)
|
||||
from vllm.model_executor.models.qwen3_5 import (
|
||||
@@ -641,6 +643,7 @@ class Qwen4ExpForCausalLM(
|
||||
SupportsPP,
|
||||
Qwen4ExpMixtureOfExperts,
|
||||
IsHybrid,
|
||||
SupportsReplaySSM,
|
||||
):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
@@ -746,20 +749,38 @@ class Qwen4ExpForCausalLM(
|
||||
num_spec=num_spec,
|
||||
)
|
||||
|
||||
@classmethod
|
||||
def _uses_gdn_recoverssm(cls, vllm_config: VllmConfig) -> bool:
|
||||
"""One predicate for the replay record's shape AND dtype (they must
|
||||
agree for the MambaSpec)."""
|
||||
num_spec = (
|
||||
vllm_config.speculative_config.num_speculative_tokens
|
||||
if vllm_config.speculative_config
|
||||
else 0
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
|
||||
|
||||
return uses_recoverssm(vllm_config.cache_config, num_spec or 0)
|
||||
|
||||
@classmethod
|
||||
def get_gdn_mamba_state_dtype_from_config(
|
||||
cls, vllm_config: VllmConfig
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
return MambaStateDtypeCalculator.gated_delta_net_state_dtype(
|
||||
vllm_config.model_config.dtype,
|
||||
vllm_config.cache_config.mamba_cache_dtype,
|
||||
vllm_config.cache_config.mamba_ssm_cache_dtype,
|
||||
) -> tuple[torch.dtype, ...]:
|
||||
dtypes: tuple[torch.dtype, ...] = (
|
||||
MambaStateDtypeCalculator.gated_delta_net_state_dtype(
|
||||
vllm_config.model_config.dtype,
|
||||
vllm_config.cache_config.mamba_cache_dtype,
|
||||
vllm_config.cache_config.mamba_ssm_cache_dtype,
|
||||
)
|
||||
)
|
||||
if cls._uses_gdn_recoverssm(vllm_config): # replay record, fp32
|
||||
dtypes = (*dtypes, torch.float32)
|
||||
return dtypes
|
||||
|
||||
@classmethod
|
||||
def get_gdn_mamba_state_shape_from_config(
|
||||
cls, vllm_config: VllmConfig
|
||||
) -> tuple[tuple[int, int], tuple[int, int]]:
|
||||
) -> MambaStateShapes:
|
||||
parallel_config = vllm_config.parallel_config
|
||||
hf_config = vllm_config.model_config.hf_text_config
|
||||
tp_size = parallel_config.tensor_parallel_size
|
||||
@@ -768,7 +789,7 @@ class Qwen4ExpForCausalLM(
|
||||
if vllm_config.speculative_config
|
||||
else 0
|
||||
)
|
||||
return MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
shapes = MambaStateShapeCalculator.gated_delta_net_state_shape(
|
||||
tp_size,
|
||||
hf_config.linear_num_key_heads,
|
||||
hf_config.linear_num_value_heads,
|
||||
@@ -777,18 +798,30 @@ class Qwen4ExpForCausalLM(
|
||||
hf_config.linear_conv_kernel_dim,
|
||||
num_spec,
|
||||
)
|
||||
if cls._uses_gdn_recoverssm(
|
||||
vllm_config
|
||||
): # replay record [HV, spec_query_len, V + K + 1]
|
||||
shapes = (
|
||||
*shapes,
|
||||
(
|
||||
hf_config.linear_num_value_heads // tp_size,
|
||||
num_spec + 1,
|
||||
hf_config.linear_value_head_dim + hf_config.linear_key_head_dim + 1,
|
||||
),
|
||||
)
|
||||
return shapes
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_dtype_from_config(
|
||||
cls,
|
||||
vllm_config: VllmConfig,
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
) -> tuple[torch.dtype, ...]:
|
||||
return cls.get_gdn_mamba_state_dtype_from_config(vllm_config)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_shape_from_config(
|
||||
cls, vllm_config: VllmConfig
|
||||
) -> tuple[tuple[int, int], tuple[int, int]]:
|
||||
) -> MambaStateShapes:
|
||||
return cls.get_gdn_mamba_state_shape_from_config(vllm_config)
|
||||
|
||||
@classmethod
|
||||
@@ -876,6 +909,7 @@ class Qwen4ExpForConditionalGeneration(
|
||||
Qwen3_5ForConditionalGeneration,
|
||||
HasInnerState,
|
||||
Qwen4ExpMixtureOfExperts,
|
||||
SupportsReplaySSM,
|
||||
):
|
||||
"""Qwen3-VL vision tower backed by the Qwen4Exp language model."""
|
||||
|
||||
@@ -1060,14 +1094,15 @@ class Qwen4ExpForConditionalGeneration(
|
||||
def get_mamba_state_dtype_from_config(
|
||||
cls,
|
||||
vllm_config: VllmConfig,
|
||||
) -> tuple[torch.dtype, torch.dtype]:
|
||||
) -> tuple[torch.dtype, ...]:
|
||||
return Qwen4ExpForCausalLM.get_mamba_state_dtype_from_config(vllm_config)
|
||||
|
||||
@classmethod
|
||||
def get_mamba_state_shape_from_config(
|
||||
def get_mamba_state_shape_from_config( # type: ignore[override]
|
||||
cls,
|
||||
vllm_config: VllmConfig,
|
||||
) -> tuple[tuple[int, int], tuple[int, int]]:
|
||||
) -> MambaStateShapes:
|
||||
# GDN RecoverSSM adds a third state (the replay record) to Qwen3.5's two.
|
||||
return Qwen4ExpForCausalLM.get_mamba_state_shape_from_config(vllm_config)
|
||||
|
||||
@classmethod
|
||||
|
||||
@@ -32,7 +32,7 @@ from vllm.model_executor.layers.vocab_parallel_embedding import (
|
||||
VocabParallelEmbedding,
|
||||
)
|
||||
from vllm.model_executor.model_loader.utils import configure_quant_config
|
||||
from vllm.model_executor.models.interfaces import SupportsPP
|
||||
from vllm.model_executor.models.interfaces import SupportsPP, SupportsReplaySSM
|
||||
from vllm.model_executor.models.qwen3_5 import Qwen3_5Model
|
||||
from vllm.model_executor.models.utils import (
|
||||
AutoWeightsLoader,
|
||||
@@ -359,7 +359,7 @@ class Qwen4ExpMultiTokenPredictor(nn.Module):
|
||||
return loader.load_weights(weights, mapper=mapper)
|
||||
|
||||
|
||||
class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts):
|
||||
class Qwen4ExpMTP(nn.Module, SupportsPP, Qwen4ExpMixtureOfExperts, SupportsReplaySSM):
|
||||
packed_modules_mapping = {
|
||||
"qkv_proj": ["q_proj", "k_proj", "v_proj"],
|
||||
"gate_up_proj": ["gate_proj", "up_proj"],
|
||||
|
||||
@@ -175,6 +175,14 @@ class Qwen4ExpPLELayer(nn.Module, MambaBase):
|
||||
return True
|
||||
|
||||
def get_attn_backend(self) -> type[PleShortConvAttentionBackend]:
|
||||
from vllm.model_executor.layers.mamba.recoverssm_utils import uses_recoverssm
|
||||
|
||||
if uses_recoverssm(self.cache_config, self.num_spec_tokens):
|
||||
from vllm.v1.attention.backends.ple_recoverssm import (
|
||||
PleRecoverSSMAttentionBackend,
|
||||
)
|
||||
|
||||
return PleRecoverSSMAttentionBackend
|
||||
return PleShortConvAttentionBackend
|
||||
|
||||
def get_state_dtype(self) -> tuple[torch.dtype, ...]:
|
||||
|
||||
@@ -0,0 +1,181 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""GDN attention metadata for RecoverSSM speculative decode (Qwen-style Gated DeltaNet).
|
||||
|
||||
Wraps the generic GDN builder: every active decode row runs through the spec path (a
|
||||
draft-less step is a T=1 window), the spec state indices keep only the checkpoint
|
||||
column, num_accepted_tokens is always 1 (the checkpoint and the compacted conv window
|
||||
already hold the accepted state), and the metadata carries what the post-sampling commit
|
||||
needs. Mirrors the Kimi-K3 KDA builder's RecoverSSM handling."""
|
||||
|
||||
from dataclasses import dataclass, field, fields
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.v1.attention.backend import AttentionCGSupport, CommonAttentionMetadata
|
||||
from vllm.v1.attention.backends.gdn_attn import (
|
||||
GDNAttentionBackend,
|
||||
GDNAttentionMetadata,
|
||||
GDNAttentionMetadataBuilder,
|
||||
)
|
||||
from vllm.v1.attention.backends.recoverssm_metadata import (
|
||||
RecoverSSMMetadata,
|
||||
RecoverSSMPostprocessMetadata,
|
||||
recoverssm_request_indices,
|
||||
recoverssm_spec_rows,
|
||||
)
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
@dataclass
|
||||
class GDNRecoverSSMCommitMetadata:
|
||||
state_indices: torch.Tensor # [num_spec_decodes, 1]
|
||||
query_start_loc: torch.Tensor # [num_spec_decodes + 1]
|
||||
request_indices: (
|
||||
torch.Tensor | None
|
||||
) # batch rows of the spec requests (None: rows 0..n-1)
|
||||
block_table: torch.Tensor | None # align mode
|
||||
num_computed_tokens: torch.Tensor | None
|
||||
block_size: int | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class GDNRecoverSSMMetadata(GDNAttentionMetadata, RecoverSSMMetadata):
|
||||
recoverssm_commit: GDNRecoverSSMCommitMetadata | None = None
|
||||
recoverssm_context: Any = field(default=None, repr=False, compare=False)
|
||||
|
||||
def commit_recoverssm_state(
|
||||
self, num_accepted_tokens: torch.Tensor
|
||||
) -> RecoverSSMPostprocessMetadata | None:
|
||||
c = self.recoverssm_commit
|
||||
if c is None:
|
||||
return None
|
||||
n = self.num_spec_decodes
|
||||
self.recoverssm_context.commit(
|
||||
num_accepted_tokens,
|
||||
c.state_indices[:n, 0],
|
||||
c.query_start_loc[: n + 1],
|
||||
request_indices=c.request_indices,
|
||||
block_table=c.block_table,
|
||||
num_computed_tokens=c.num_computed_tokens,
|
||||
mamba_block_size=c.block_size,
|
||||
)
|
||||
if c.block_table is None:
|
||||
return None
|
||||
assert c.block_size is not None and c.num_computed_tokens is not None
|
||||
return RecoverSSMPostprocessMetadata(
|
||||
num_spec_decodes=n,
|
||||
request_indices=c.request_indices,
|
||||
block_table=c.block_table,
|
||||
num_computed_tokens=c.num_computed_tokens,
|
||||
block_size=c.block_size,
|
||||
)
|
||||
|
||||
|
||||
class GDNRecoverSSMMetadataBuilder(GDNAttentionMetadataBuilder):
|
||||
# The verify/commit metadata is not staged for FULL graph replay yet; declaring
|
||||
# NEVER makes the runner fall back to PIECEWISE (attention compiled piecewise).
|
||||
_cudagraph_support = AttentionCGSupport.NEVER
|
||||
|
||||
def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None:
|
||||
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
|
||||
# The builder may be created before the runner resolves the CUDA graph mode.
|
||||
self.use_full_cuda_graph = False
|
||||
max_reqs = vllm_config.scheduler_config.max_num_seqs
|
||||
self._ones = torch.ones(max_reqs, dtype=torch.int32, device=device)
|
||||
self._context: Any = None
|
||||
self._logged = False
|
||||
|
||||
def _get_context(self):
|
||||
if self._context is None:
|
||||
from vllm.model_executor.layers.mamba.gdn.recoverssm_gdn import (
|
||||
GDNRecoverSSMCommitContext,
|
||||
)
|
||||
|
||||
fc = self.vllm_config.compilation_config.static_forward_context
|
||||
layers = [fc[name] for name in self.layer_names]
|
||||
self._context = GDNRecoverSSMCommitContext.create(
|
||||
layers,
|
||||
spec_query_len=1 + self.num_spec,
|
||||
max_num_reqs=self.vllm_config.scheduler_config.max_num_seqs,
|
||||
)
|
||||
logger.info(
|
||||
"GDN RecoverSSM commit context: %d layers, spec_query_len %d",
|
||||
len(layers),
|
||||
1 + self.num_spec,
|
||||
)
|
||||
return self._context
|
||||
|
||||
def build( # type: ignore[override]
|
||||
self,
|
||||
common_prefix_len: int,
|
||||
common_attn_metadata: CommonAttentionMetadata,
|
||||
num_accepted_tokens: torch.Tensor | None = None,
|
||||
num_decode_draft_tokens_cpu: torch.Tensor | None = None,
|
||||
fast_build: bool = False,
|
||||
) -> "GDNRecoverSSMMetadata":
|
||||
m = common_attn_metadata
|
||||
rows = None
|
||||
if self.use_spec_decode and num_decode_draft_tokens_cpu is not None:
|
||||
if m.is_prefilling is None or m.is_prefilling.device.type != "cpu":
|
||||
raise ValueError("GDN RecoverSSM needs the CPU is_prefilling mask")
|
||||
# Positive counts only defeat the base builder's "no drafts -> regular
|
||||
# decode" shortcut: RecoverSSM keeps every post-prefill row on the spec path
|
||||
# (extended conv window).
|
||||
num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows(
|
||||
m.is_prefilling,
|
||||
m.query_start_loc_cpu,
|
||||
num_decode_draft_tokens_cpu,
|
||||
m.num_reqs,
|
||||
)
|
||||
meta = super().build(
|
||||
common_prefix_len,
|
||||
m,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
|
||||
fast_build=fast_build,
|
||||
)
|
||||
base = {f.name: getattr(meta, f.name) for f in fields(meta)}
|
||||
commit = None
|
||||
if meta.num_spec_decodes > 0:
|
||||
n = meta.num_spec_decodes
|
||||
assert meta.spec_state_indices_tensor is not None
|
||||
base["spec_state_indices_tensor"] = meta.spec_state_indices_tensor[:, :1]
|
||||
base["num_accepted_tokens"] = self._ones[:n]
|
||||
if rows is None:
|
||||
raise ValueError("GDN RecoverSSM: spec rows without draft counts")
|
||||
request_indices = recoverssm_request_indices(
|
||||
rows, n, m.query_start_loc.device
|
||||
)
|
||||
align = self.vllm_config.cache_config.mamba_cache_mode == "align"
|
||||
commit = GDNRecoverSSMCommitMetadata(
|
||||
state_indices=base["spec_state_indices_tensor"],
|
||||
query_start_loc=meta.spec_query_start_loc,
|
||||
request_indices=request_indices,
|
||||
block_table=m.block_table_tensor if align else None,
|
||||
num_computed_tokens=m.compute_num_computed_tokens() if align else None,
|
||||
block_size=self.kv_cache_spec.block_size if align else None,
|
||||
)
|
||||
if not self._logged:
|
||||
self._logged = True
|
||||
logger.info(
|
||||
"GDN RecoverSSM path taken: %d spec rows, align=%s", n, align
|
||||
)
|
||||
return GDNRecoverSSMMetadata(
|
||||
**base,
|
||||
recoverssm_commit=commit,
|
||||
recoverssm_context=self._get_context() if commit is not None else None,
|
||||
)
|
||||
|
||||
|
||||
class GDNRecoverSSMAttentionBackend(GDNAttentionBackend):
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "GDN_RECOVERSSM"
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type[GDNRecoverSSMMetadataBuilder]:
|
||||
return GDNRecoverSSMMetadataBuilder
|
||||
@@ -112,7 +112,12 @@ class BaseMambaAttentionMetadataBuilder(AttentionMetadataBuilder[M], abc.ABC):
|
||||
self.compilation_config = vllm_config.compilation_config
|
||||
self.num_spec_tokens: int = vllm_config.num_speculative_tokens
|
||||
self.use_spec_decode = self.num_spec_tokens > 0
|
||||
self.use_replayssm = vllm_config.cache_config.use_replayssm
|
||||
# RecoverSSM also rides on --use-replayssm but keeps its own metadata; the
|
||||
# Mamba2 ReplaySSM ring-buffer workspaces below are not used then.
|
||||
self.use_replayssm = (
|
||||
vllm_config.cache_config.use_replayssm
|
||||
and not vllm_config.cache_config.use_recoverssm
|
||||
)
|
||||
self.replayssm_buffer_len = vllm_config.cache_config.replayssm_buffer_len
|
||||
self.use_flashinfer_replayssm = (
|
||||
self.use_replayssm
|
||||
|
||||
@@ -0,0 +1,326 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
||||
"""RecoverSSM protocol for the Qwen4Exp PLE short-conv layer.
|
||||
|
||||
When the GDN layers run RecoverSSM, the runner resets the shared num_accepted_tokens to
|
||||
1 after each commit (align mode). The PLE dilated short conv keeps an extended window in
|
||||
one state block and reads it at num_accepted - 1, so it must follow the same protocol:
|
||||
every active decode row takes the spec path, the conv reads at offset 0, and after
|
||||
sampling the accepted window is compacted to the front (the shared RecoverSSM compaction
|
||||
kernel; in align mode it also writes the block-boundary window)."""
|
||||
|
||||
from dataclasses import dataclass, field, fields
|
||||
from typing import Any
|
||||
|
||||
import torch
|
||||
|
||||
from vllm.logger import init_logger
|
||||
from vllm.model_executor.layers.mamba.ops.recoverssm import (
|
||||
compact_conv_state_kernel,
|
||||
prepare_commit_plan_kernel,
|
||||
)
|
||||
from vllm.model_executor.layers.mamba.recoverssm_utils import recoverssm_require
|
||||
from vllm.triton_utils import triton
|
||||
from vllm.v1.attention.backend import AttentionCGSupport
|
||||
from vllm.v1.attention.backends.recoverssm_metadata import (
|
||||
RecoverSSMMetadata,
|
||||
RecoverSSMPostprocessMetadata,
|
||||
recoverssm_request_indices,
|
||||
recoverssm_spec_rows,
|
||||
)
|
||||
from vllm.v1.attention.backends.short_conv_attn import (
|
||||
PleShortConvAttentionBackend,
|
||||
PleShortConvAttentionMetadata,
|
||||
PleShortConvAttentionMetadataBuilder,
|
||||
)
|
||||
from vllm.v1.attention.backends.utils import NULL_BLOCK_ID
|
||||
|
||||
logger = init_logger(__name__)
|
||||
|
||||
|
||||
def _require(cond: bool, msg: str) -> None:
|
||||
recoverssm_require(cond, msg, "PLE RecoverSSM")
|
||||
|
||||
|
||||
class _PleConvCommit:
|
||||
def __init__(self, conv_states, spec_query_len: int, max_num_reqs: int):
|
||||
from vllm.model_executor.layers.mamba.mamba_utils import is_conv_state_dim_first
|
||||
|
||||
_require(len(conv_states) > 0, "PLE commit requires at least one layer")
|
||||
if not is_conv_state_dim_first():
|
||||
conv_states = [
|
||||
s.transpose(-1, -2) for s in conv_states
|
||||
] # -> [blocks, C, W]
|
||||
self.conv_states = list(conv_states)
|
||||
ref = self.conv_states[0]
|
||||
_require(ref.ndim == 3, "PLE conv state must be [blocks, dim, window]")
|
||||
dev = ref.device
|
||||
for s in self.conv_states:
|
||||
_require(
|
||||
s.shape == ref.shape
|
||||
and s.dtype == ref.dtype
|
||||
and s.stride()[1:] == ref.stride()[1:]
|
||||
and s.device == dev,
|
||||
"PLE layers need matching conv states",
|
||||
)
|
||||
self.conv_dim, conv_len = ref.shape[1], ref.shape[2]
|
||||
self.hist = conv_len - spec_query_len + 1
|
||||
_require(self.hist > 0, "PLE conv state is shorter than its window")
|
||||
self.spec_query_len = spec_query_len
|
||||
t = lambda xs: torch.tensor(xs, dtype=torch.int64, device=dev)
|
||||
self.base = t([s.data_ptr() for s in self.conv_states])
|
||||
self.bstride = t([s.stride(0) for s in self.conv_states])
|
||||
self.dstride = t([s.stride(1) for s in self.conv_states])
|
||||
self.tstride = t([s.stride(2) for s in self.conv_states])
|
||||
z = lambda: torch.empty(max_num_reqs, dtype=torch.int32, device=dev)
|
||||
self.commit_lens, self.final_idx, self.bnd_idx, self.bnd_len = (
|
||||
z(),
|
||||
z(),
|
||||
z(),
|
||||
z(),
|
||||
)
|
||||
|
||||
def commit(
|
||||
self,
|
||||
num_accepted_tokens,
|
||||
state_indices,
|
||||
query_start_loc,
|
||||
request_indices=None,
|
||||
block_table=None,
|
||||
num_computed_tokens=None,
|
||||
mamba_block_size=None,
|
||||
):
|
||||
batch = state_indices.shape[0]
|
||||
if batch == 0:
|
||||
return
|
||||
_require(state_indices.ndim == 1, "state indices must be one-dimensional")
|
||||
_require(
|
||||
batch <= self.commit_lens.shape[0],
|
||||
"PLE commit batch exceeds its plan capacity",
|
||||
)
|
||||
_require(
|
||||
query_start_loc.ndim == 1 and query_start_loc.shape[0] == batch + 1,
|
||||
"PLE commit metadata is incompatible",
|
||||
)
|
||||
_require(
|
||||
request_indices is None or request_indices.shape[0] >= batch,
|
||||
"PLE request mapping is too short",
|
||||
)
|
||||
_require(
|
||||
num_accepted_tokens.ndim == 1
|
||||
and (request_indices is not None or num_accepted_tokens.shape[0] >= batch),
|
||||
"PLE accepted-token counts are too short",
|
||||
)
|
||||
_align = (block_table, num_computed_tokens, mamba_block_size)
|
||||
_require(
|
||||
all(x is None for x in _align) or all(x is not None for x in _align),
|
||||
"PLE align metadata is incomplete",
|
||||
)
|
||||
_require(
|
||||
mamba_block_size is None or mamba_block_size >= self.spec_query_len,
|
||||
"PLE align block size must cover one speculative window",
|
||||
)
|
||||
_require(
|
||||
block_table is None or block_table.ndim == 2,
|
||||
"PLE block table must be two-dimensional",
|
||||
)
|
||||
_dev = self.conv_states[0].device
|
||||
_require(
|
||||
all(
|
||||
t_ is None or t_.device == _dev
|
||||
for t_ in (
|
||||
num_accepted_tokens,
|
||||
state_indices,
|
||||
query_start_loc,
|
||||
request_indices,
|
||||
block_table,
|
||||
num_computed_tokens,
|
||||
)
|
||||
),
|
||||
"PLE commit inputs must be on the same device",
|
||||
)
|
||||
bt = (0, 0) if block_table is None else block_table.stride()
|
||||
prepare_commit_plan_kernel[(batch,)](
|
||||
num_accepted_tokens,
|
||||
request_indices,
|
||||
state_indices,
|
||||
query_start_loc,
|
||||
block_table,
|
||||
num_computed_tokens,
|
||||
self.commit_lens,
|
||||
self.final_idx,
|
||||
self.bnd_idx,
|
||||
self.bnd_len,
|
||||
NULL_BLOCK_ID,
|
||||
mamba_block_size or 1,
|
||||
block_table.shape[1] if block_table is not None else 1,
|
||||
num_accepted_tokens.stride(0),
|
||||
request_indices.stride(0) if request_indices is not None else 0,
|
||||
state_indices.stride(0),
|
||||
query_start_loc.stride(0),
|
||||
bt[0],
|
||||
bt[1],
|
||||
0 if num_computed_tokens is None else num_computed_tokens.stride(0),
|
||||
SPEC_QUERY_LEN=self.spec_query_len,
|
||||
num_warps=1,
|
||||
)
|
||||
compact_conv_state_kernel[
|
||||
(triton.cdiv(self.conv_dim, 256), batch, len(self.conv_states))
|
||||
](
|
||||
self.conv_states[0],
|
||||
self.base,
|
||||
self.bstride,
|
||||
self.dstride,
|
||||
self.tstride,
|
||||
state_indices,
|
||||
self.commit_lens,
|
||||
self.final_idx,
|
||||
self.bnd_idx,
|
||||
self.bnd_len,
|
||||
NULL_BLOCK_ID,
|
||||
self.conv_dim,
|
||||
self.hist,
|
||||
state_indices.stride(0),
|
||||
BLOCK_D=256,
|
||||
BLOCK_HISTORY=triton.next_power_of_2(self.hist),
|
||||
ALIGN_MODE=block_table is not None,
|
||||
num_warps=4,
|
||||
)
|
||||
|
||||
|
||||
@dataclass
|
||||
class PleRecoverSSMCommitMetadata:
|
||||
state_indices: torch.Tensor # [n]
|
||||
query_start_loc: torch.Tensor # [n + 1]
|
||||
request_indices: torch.Tensor | None
|
||||
block_table: torch.Tensor | None
|
||||
num_computed_tokens: torch.Tensor | None
|
||||
block_size: int | None
|
||||
|
||||
|
||||
@dataclass
|
||||
class PleRecoverSSMMetadata(PleShortConvAttentionMetadata, RecoverSSMMetadata):
|
||||
recoverssm_commit: PleRecoverSSMCommitMetadata | None = None
|
||||
recoverssm_context: Any = field(default=None, repr=False, compare=False)
|
||||
|
||||
def commit_recoverssm_state(self, num_accepted_tokens):
|
||||
c = self.recoverssm_commit
|
||||
if c is None:
|
||||
return None
|
||||
n = self.num_spec_decodes
|
||||
self.recoverssm_context.commit(
|
||||
num_accepted_tokens,
|
||||
c.state_indices[:n],
|
||||
c.query_start_loc[: n + 1],
|
||||
request_indices=c.request_indices,
|
||||
block_table=c.block_table,
|
||||
num_computed_tokens=c.num_computed_tokens,
|
||||
mamba_block_size=c.block_size,
|
||||
)
|
||||
if c.block_table is None:
|
||||
return None
|
||||
assert c.block_size is not None and c.num_computed_tokens is not None
|
||||
return RecoverSSMPostprocessMetadata(
|
||||
num_spec_decodes=n,
|
||||
request_indices=c.request_indices,
|
||||
block_table=c.block_table,
|
||||
num_computed_tokens=c.num_computed_tokens,
|
||||
block_size=c.block_size,
|
||||
)
|
||||
|
||||
|
||||
class PleRecoverSSMMetadataBuilder(PleShortConvAttentionMetadataBuilder):
|
||||
# Same as the GDN RecoverSSM builder: PIECEWISE only for now.
|
||||
_cudagraph_support = AttentionCGSupport.NEVER
|
||||
|
||||
def __init__(self, kv_cache_spec, layer_names, vllm_config, device) -> None:
|
||||
super().__init__(kv_cache_spec, layer_names, vllm_config, device)
|
||||
self.use_full_cuda_graph = False
|
||||
self._ones = torch.ones(
|
||||
vllm_config.scheduler_config.max_num_seqs, dtype=torch.int32, device=device
|
||||
)
|
||||
self._ctx: _PleConvCommit | None = None
|
||||
self._logged = False
|
||||
|
||||
def _get_ctx(self):
|
||||
if self._ctx is None:
|
||||
fc = self.vllm_config.compilation_config.static_forward_context
|
||||
convs = [fc[name].kv_cache[0] for name in self.layer_names]
|
||||
ctx = _PleConvCommit(
|
||||
convs, 1 + self.num_spec, self.vllm_config.scheduler_config.max_num_seqs
|
||||
)
|
||||
self._ctx = ctx
|
||||
logger.info(
|
||||
"PLE RecoverSSM commit context: %d layers, history %d",
|
||||
len(convs),
|
||||
ctx.hist,
|
||||
)
|
||||
return self._ctx
|
||||
|
||||
def build(
|
||||
self,
|
||||
common_prefix_len,
|
||||
common_attn_metadata,
|
||||
fast_build=False,
|
||||
*,
|
||||
num_accepted_tokens=None,
|
||||
num_decode_draft_tokens_cpu=None,
|
||||
**kwargs,
|
||||
): # type: ignore[override]
|
||||
m = common_attn_metadata
|
||||
rows = None
|
||||
if self.use_spec_decode and num_decode_draft_tokens_cpu is not None:
|
||||
if m.is_prefilling is None or m.is_prefilling.device.type != "cpu":
|
||||
raise ValueError("PLE RecoverSSM needs the CPU is_prefilling mask")
|
||||
# the same classification as the GDN builder, so both commits map the same
|
||||
# requests
|
||||
num_decode_draft_tokens_cpu, rows = recoverssm_spec_rows(
|
||||
m.is_prefilling,
|
||||
m.query_start_loc_cpu,
|
||||
num_decode_draft_tokens_cpu,
|
||||
m.num_reqs,
|
||||
)
|
||||
meta = super().build(
|
||||
common_prefix_len,
|
||||
m,
|
||||
fast_build,
|
||||
num_accepted_tokens=num_accepted_tokens,
|
||||
num_decode_draft_tokens_cpu=num_decode_draft_tokens_cpu,
|
||||
**kwargs,
|
||||
)
|
||||
if (
|
||||
not isinstance(meta, PleShortConvAttentionMetadata)
|
||||
or meta.num_spec_decodes == 0
|
||||
):
|
||||
return meta
|
||||
base = {f.name: getattr(meta, f.name) for f in fields(meta)}
|
||||
n = meta.num_spec_decodes
|
||||
base["num_accepted_tokens"] = self._ones[:n]
|
||||
if rows is None:
|
||||
raise ValueError("PLE RecoverSSM: spec rows without draft counts")
|
||||
req_idx = recoverssm_request_indices(rows, n, m.query_start_loc.device)
|
||||
align = self.vllm_config.cache_config.mamba_cache_mode == "align"
|
||||
commit = PleRecoverSSMCommitMetadata(
|
||||
state_indices=meta.spec_state_indices_tensor,
|
||||
query_start_loc=meta.spec_query_start_loc,
|
||||
request_indices=req_idx,
|
||||
block_table=m.block_table_tensor if align else None,
|
||||
num_computed_tokens=m.compute_num_computed_tokens() if align else None,
|
||||
block_size=self.kv_cache_spec.block_size if align else None,
|
||||
)
|
||||
if not self._logged:
|
||||
self._logged = True
|
||||
logger.info("PLE RecoverSSM path taken: %d spec rows, align=%s", n, align)
|
||||
return PleRecoverSSMMetadata(
|
||||
**base, recoverssm_commit=commit, recoverssm_context=self._get_ctx()
|
||||
)
|
||||
|
||||
|
||||
class PleRecoverSSMAttentionBackend(PleShortConvAttentionBackend):
|
||||
@staticmethod
|
||||
def get_name() -> str:
|
||||
return "PLE_SHORT_CONV_RECOVERSSM"
|
||||
|
||||
@staticmethod
|
||||
def get_builder_cls() -> type[PleRecoverSSMMetadataBuilder]:
|
||||
return PleRecoverSSMMetadataBuilder
|
||||
Reference in New Issue
Block a user