feat(qwen4_exp): support Qwen3.8-Flash-Next (#257)
Serve Qwen3.8-Flash-Next (HF model_type qwen4_exp) text-only: 36 GDN + 12 QSA compressed-sparse attention layers on 4 hyper-connection residual streams, a PLE n-gram embedding layer backed by a 47.7 GiB pinned-host table with UVA gather, and 512 NVFP4 / block-fp8 routed experts (top-10) plus a gated shared expert. - attention: qsa_sparse backend (AttnType.QSA) over QSAKVCache -- paged GQA K/V, a 1/ratio compressed index-key slab shadowing the KV pages, and a per-request pending ring sized from index_ratio - kvcache: declarative slot-sibling states (ModelConfig.slot_states) on LinearStatePool carry the PLE conv history and n-gram context through the hybrid-radix snapshot/COW lifecycle - scheduler: hybrid prefill chunks align to the page size so snapshots land on donatable boundaries - kernels: triton kernels adapted from vLLM/SGLang (hc, qsa, ple gather, moe router / shared gate) plus an original radix block top-k; int64 row addressing throughout - moe: non-power-of-2 top-k router, deep-K marlin decode config, one fp8 scale-bank padding rule shared with the AOT row table - engine: the PLE table load reserves its pinned bytes from the pin budget before the expert banks plan their residency
This commit is contained in:
@@ -37,7 +37,7 @@ UID = 2
|
||||
def _pool(num_slots=16):
|
||||
g = LinearGatedDeltaGroupConfig(
|
||||
name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4,
|
||||
key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate=True,
|
||||
key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu",
|
||||
)
|
||||
return LinearStatePool(group=g, num_slots=num_slots, dtype=torch.bfloat16,
|
||||
device=torch.device("cpu"), tp_size=1)
|
||||
|
||||
@@ -16,7 +16,7 @@ from freetoken.scheduler.cache import CacheManager
|
||||
def _pool(num_slots=16):
|
||||
g = LinearGatedDeltaGroupConfig(
|
||||
name="linear", layer_ids=(0,), num_key_heads=2, num_value_heads=4,
|
||||
key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate=True,
|
||||
key_head_dim=16, value_head_dim=16, conv_kernel_dim=4, output_gate="silu",
|
||||
)
|
||||
return LinearStatePool(group=g, num_slots=num_slots, dtype=torch.bfloat16,
|
||||
device=torch.device("cpu"), tp_size=1)
|
||||
@@ -116,6 +116,43 @@ def test_rebuild_reclaims_donated_gdn_slots():
|
||||
assert pool.num_free_slots == pool.num_slots - 1 # all GDN slots reclaimed (no leak)
|
||||
|
||||
|
||||
def test_prefill_chunk_ends_on_a_page_boundary():
|
||||
"""A hybrid chunk must end page-aligned: the snapshot commit skips any other boundary."""
|
||||
from freetoken.scheduler.prefill import ChunkedReq, PrefillAdder
|
||||
from freetoken.scheduler.table import TableManager
|
||||
from freetoken.scheduler.utils import PendingReq
|
||||
|
||||
pool = _pool()
|
||||
pt = torch.zeros(4, 512, dtype=torch.int32)
|
||||
cm = CacheManager(64, 64, pt, "hybrid_radix", linear_state_pool=pool)
|
||||
assert cm.prefill_chunk_align == 64
|
||||
tm = TableManager(max_running_reqs=4, page_table=pt)
|
||||
pending = PendingReq(0, torch.arange(300, dtype=torch.int32), SamplingParams(max_tokens=1))
|
||||
|
||||
adder = PrefillAdder(token_budget=100, reserved_size=0, cache_manager=cm, table_manager=tm)
|
||||
req = adder.try_add_one(pending)
|
||||
assert isinstance(req, ChunkedReq) and req.extend_len == 64
|
||||
|
||||
# a budget below one page keeps the unaligned chunk rather than stalling the request
|
||||
adder = PrefillAdder(token_budget=40, reserved_size=0, cache_manager=cm, table_manager=tm)
|
||||
assert adder.try_add_one(pending).extend_len == 40
|
||||
|
||||
|
||||
def test_naive_cache_does_not_align_prefill_chunks():
|
||||
"""The alignment hook is hybrid-only; every other cache keeps the raw budget chunk."""
|
||||
from freetoken.scheduler.prefill import PrefillAdder
|
||||
from freetoken.scheduler.table import TableManager
|
||||
from freetoken.scheduler.utils import PendingReq
|
||||
|
||||
pt = torch.zeros(4, 512, dtype=torch.int32)
|
||||
cm = CacheManager(64, 64, pt, "radix")
|
||||
assert cm.prefill_chunk_align == 1
|
||||
tm = TableManager(max_running_reqs=4, page_table=pt)
|
||||
adder = PrefillAdder(token_budget=100, reserved_size=0, cache_manager=cm, table_manager=tm)
|
||||
pending = PendingReq(0, torch.arange(300, dtype=torch.int32), SamplingParams(max_tokens=1))
|
||||
assert adder.try_add_one(pending).extend_len == 100
|
||||
|
||||
|
||||
def test_pool_sizing_covers_4mr_floor():
|
||||
"""C6: pool must reserve the 4-slot-per-request non-evictable floor even at a tiny ratio."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
Reference in New Issue
Block a user