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:
Xiaoze Fan
2026-08-28 15:33:15 -07:00
committed by GitHub
parent 9ef3651309
commit bd8f3d519a
66 changed files with 10060 additions and 106 deletions
@@ -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)
+38 -1
View File
@@ -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