[MLX] Window-bounded SWA KV storage and in-graph sampling (#34166)
Co-authored-by: Siming Deng <siming_deng_stat@163.com> Co-authored-by: R0CKSTAR <yeahdongcn@gmail.com> Co-authored-by: Jiminator <Jiminator@users.noreply.github.com> Co-authored-by: damahua <damahua@users.noreply.github.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Alex Nails
co-authored by
Siming Deng
R0CKSTAR
Jiminator
damahua
Claude Opus 5
Siming Deng
R0CKSTAR
Jiminator
damahua
Claude Opus 5
parent
449f0da78f
commit
2969ab3d41
@@ -5,8 +5,9 @@ uses per-head attention sinks, so it exercises the MLX backend's
|
||||
sliding-window path end to end. Two guards:
|
||||
|
||||
1. ``TestGptOssMlxCorrectness`` — black-box serving smoke against a running
|
||||
server, including a >128-token prompt so the sliding window actually
|
||||
engages.
|
||||
server with the radix cache enabled (the default KV path), including a
|
||||
>128-token prompt so the sliding window actually engages and a repeated
|
||||
prompt so a radix prefix hit must reproduce the cold greedy output.
|
||||
2. ``TestGptOssMlxReferenceCorrectness`` — token-for-token equivalence of
|
||||
``MlxModelRunner`` greedy decoding against raw, unpatched mlx_lm greedy
|
||||
generation. SGLang keeps full KV and applies banded masks /
|
||||
@@ -117,10 +118,14 @@ class TestGptOssMlxCorrectness(CustomTestCase):
|
||||
cls.base_url,
|
||||
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
|
||||
other_args=[
|
||||
# Radix cache stays enabled (the default): sliding-window
|
||||
# layers keep windowed per-request KV, the shared pool holds
|
||||
# full-attention layers, and prefix hits recompute the
|
||||
# prefix, so serving must stay correct without
|
||||
# --disable-radix-cache.
|
||||
"--trust-remote-code",
|
||||
"--tp-size",
|
||||
"1",
|
||||
"--disable-radix-cache",
|
||||
"--disable-cuda-graph",
|
||||
"--mem-fraction-static",
|
||||
MEM_FRACTION_STATIC,
|
||||
@@ -189,6 +194,24 @@ class TestGptOssMlxCorrectness(CustomTestCase):
|
||||
)
|
||||
self.assertIn("BLUEBERRY", text.upper())
|
||||
|
||||
def test_radix_prefix_hit_reproduces_greedy_output(self):
|
||||
# The server runs with the radix cache enabled. Sending the same
|
||||
# >128-token prompt twice makes the second request hit the cached
|
||||
# prefix; on sliding-window models the runner recomputes the prefix
|
||||
# (windowed KV keeps no pool history), and greedy output must be
|
||||
# identical to the cold request.
|
||||
messages = [
|
||||
{"role": "system", "content": "You are a concise assistant."},
|
||||
{
|
||||
"role": "user",
|
||||
"content": _NUMBER_LIST
|
||||
+ ". Which number comes right after 41? Answer briefly.",
|
||||
},
|
||||
]
|
||||
cold = self._chat(messages, max_tokens=48)
|
||||
hit = self._chat(messages, max_tokens=48)
|
||||
self.assertEqual(cold, hit)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestGptOssMlxReferenceCorrectness(CustomTestCase):
|
||||
|
||||
@@ -46,6 +46,7 @@ if _HAS_MLX:
|
||||
MlxPendingJob,
|
||||
SchedulerMlxOverlapMixin,
|
||||
)
|
||||
from sglang.srt.hardware_backend.mlx.tp_worker import MlxLaunch
|
||||
from sglang.srt.managers.scheduler_components import (
|
||||
batch_result_processor as batch_result_processor_module,
|
||||
)
|
||||
@@ -309,7 +310,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
new_slot_ids=[4],
|
||||
req_pool_idx=0,
|
||||
)
|
||||
MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache)
|
||||
runner.eval_pending(pending)
|
||||
mx.eval(*runner._attention_kv_pool.all_buffers())
|
||||
runner.prefill_finalize(pending)
|
||||
|
||||
@@ -350,7 +351,8 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
calls.append(
|
||||
(len(caches), batched_input.tolist(), list(helper_req_ids))
|
||||
)
|
||||
return mx.array(list(range(len(caches))), dtype=mx.int32)
|
||||
# Last-token logits whose argmax is the row index.
|
||||
return mx.eye(len(caches), 8, dtype=mx.float32)
|
||||
|
||||
def fail_native(*args, **kwargs):
|
||||
raise AssertionError("dense decode should use batched attention")
|
||||
@@ -386,7 +388,8 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
|
||||
def fake_batched(caches, batched_input, helper_req_ids):
|
||||
calls.append((len(caches), batched_input.tolist(), list(helper_req_ids)))
|
||||
return mx.array([8], dtype=mx.int32)
|
||||
# Last-token logits whose argmax is token 8.
|
||||
return mx.arange(9, dtype=mx.float32)[None, :]
|
||||
|
||||
def fail_native(*args, **kwargs):
|
||||
raise AssertionError("dense chained decode should use batched attention")
|
||||
@@ -502,12 +505,13 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
]
|
||||
]
|
||||
|
||||
lazy_tokens = runner._decode_with_batched_attention(
|
||||
lazy_logits = runner._decode_with_batched_attention(
|
||||
cache,
|
||||
mx.array([[7]], dtype=mx.int32),
|
||||
["r0"],
|
||||
)
|
||||
mx.eval(lazy_tokens, *MlxModelRunner._cache_state_arrays(cache))
|
||||
lazy_tokens = mx.argmax(lazy_logits, axis=-1)
|
||||
mx.eval(lazy_tokens, *MlxModelRunner.cache_state_arrays(cache))
|
||||
|
||||
self.assertEqual(lazy_tokens.tolist(), [0])
|
||||
self.assertEqual(cache[0][0].offset, 1)
|
||||
@@ -558,6 +562,9 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
req_pool_idx={"r0": 0, "r1": 1},
|
||||
req_to_token_pool=req_to_token_pool,
|
||||
attention_layer_indices=[0],
|
||||
# The fused scatter addresses pool buffers by full-attention index,
|
||||
# so the context requires the map whenever the RoPE kernel is live.
|
||||
full_kv_pool_index_by_layer={0: 0},
|
||||
)
|
||||
|
||||
self.assertEqual(ctx.seq_lens, [1, 2])
|
||||
@@ -578,7 +585,8 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
|
||||
def fake_hybrid(caches, batched_input, helper_req_ids):
|
||||
calls.append((len(caches), batched_input.tolist(), list(helper_req_ids)))
|
||||
return mx.array([4, 5], dtype=mx.int32)
|
||||
# Last-token logits whose argmax is 4 for row 0, 5 for row 1.
|
||||
return mx.eye(8, dtype=mx.float32)[4:6]
|
||||
|
||||
def fail_batched(*args, **kwargs):
|
||||
raise AssertionError(
|
||||
@@ -713,7 +721,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
new_slot_ids=[4],
|
||||
req_pool_idx=req.req_pool_idx,
|
||||
)
|
||||
MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache)
|
||||
runner.eval_pending(pending)
|
||||
runner.prefill_finalize(pending)
|
||||
|
||||
self.assertEqual(runner.model.seen_inputs, [[[13]]])
|
||||
@@ -770,7 +778,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
req_pool_idx=req.req_pool_idx,
|
||||
req=req,
|
||||
)
|
||||
MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache)
|
||||
runner.eval_pending(pending)
|
||||
runner.prefill_finalize(pending)
|
||||
tracked = [FakeNativeCache(), None]
|
||||
runner._req_to_token_pool.auxiliary_state_pool.restore_cache(
|
||||
@@ -832,7 +840,7 @@ class TestMlxAuxiliaryStateRunnerCache(unittest.TestCase):
|
||||
req_pool_idx=req.req_pool_idx,
|
||||
req=req,
|
||||
)
|
||||
MlxModelRunner._eval_with_cache(pending.lazy_token, pending.cache)
|
||||
runner.eval_pending(pending)
|
||||
runner.prefill_finalize(pending)
|
||||
tracked = [FakeNativeCache(), None]
|
||||
runner._req_to_token_pool.auxiliary_state_pool.restore_cache(
|
||||
@@ -1094,11 +1102,13 @@ class TestMlxOverlapScheduler(unittest.TestCase):
|
||||
scheduler.last_batch = stale_batch
|
||||
|
||||
pending = MlxPendingJob(
|
||||
lazy_tokens=None,
|
||||
prefills=["prefill"],
|
||||
extends=[],
|
||||
decode=None,
|
||||
mode="extend",
|
||||
launch=MlxLaunch(
|
||||
lazy_tokens=None,
|
||||
prefills=["prefill"],
|
||||
extends=[],
|
||||
decode=None,
|
||||
mode="extend",
|
||||
),
|
||||
batch_copy=batch_copy,
|
||||
schedule_batch=schedule_batch,
|
||||
reqs=[SimpleNamespace(rid="r0")],
|
||||
|
||||
@@ -203,8 +203,14 @@ class TestSchedulerProfilerManagerMPS(unittest.TestCase):
|
||||
mgr._init_profile(output_dir, None, None, None, None, None, False, "test")
|
||||
return mgr
|
||||
|
||||
# MetalCaptureProfiler has two strategies: start_mlx drives
|
||||
# mx.metal.start_capture, start_mps drives torch.mps.profiler.metal_capture.
|
||||
# This manager takes the MPS one, so that is the symbol to stand in for --
|
||||
# patching mx.metal here leaves the real Metal capture running, which fails
|
||||
# with "Capture layer is not inserted" unless MTL_CAPTURE_ENABLED=1 is set
|
||||
# in the environment.
|
||||
def test_start_profile_failure_does_not_crash(self):
|
||||
import mlx.core as mx
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.profiler import (
|
||||
apply_metal_profiler_patches,
|
||||
@@ -215,20 +221,22 @@ class TestSchedulerProfilerManagerMPS(unittest.TestCase):
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
mgr = self._make_manager(tmp)
|
||||
with patch.object(
|
||||
mx.metal,
|
||||
"start_capture",
|
||||
torch.mps.profiler,
|
||||
"metal_capture",
|
||||
side_effect=RuntimeError("Capture layer is not inserted"),
|
||||
):
|
||||
result = mgr._start_profile()
|
||||
|
||||
self.assertFalse(result.success)
|
||||
self.assertIn("Capture layer is not inserted", result.message)
|
||||
self.assertFalse(mgr.profile_in_progress)
|
||||
self.assertIsNone(mgr.torch_profiler)
|
||||
|
||||
def test_start_profile_success_with_mock_capture(self):
|
||||
from unittest.mock import MagicMock
|
||||
from unittest.mock import patch as mock_patch
|
||||
|
||||
import mlx.core as mx
|
||||
import torch
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.profiler import (
|
||||
apply_metal_profiler_patches,
|
||||
@@ -238,14 +246,17 @@ class TestSchedulerProfilerManagerMPS(unittest.TestCase):
|
||||
|
||||
with tempfile.TemporaryDirectory() as tmp:
|
||||
mgr = self._make_manager(tmp)
|
||||
with mock_patch.object(mx.metal, "start_capture"), mock_patch.object(
|
||||
mx.metal, "stop_capture"
|
||||
capture_ctx = MagicMock()
|
||||
with mock_patch.object(
|
||||
torch.mps.profiler, "metal_capture", return_value=capture_ctx
|
||||
), mock_patch("torch.distributed.barrier"):
|
||||
result = mgr._start_profile()
|
||||
self.assertTrue(result.success)
|
||||
self.assertTrue(result.success, result.message)
|
||||
self.assertTrue(mgr.profile_in_progress)
|
||||
capture_ctx.__enter__.assert_called_once()
|
||||
mgr._stop_profile()
|
||||
self.assertFalse(mgr.profile_in_progress)
|
||||
capture_ctx.__exit__.assert_called_once()
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -202,6 +202,13 @@ class TestMlxReferenceCorrectness(CustomTestCase):
|
||||
self.runner.remove_request(rid)
|
||||
return out
|
||||
|
||||
def _truncate_at_eos(self, seq):
|
||||
"""``seq`` up to and including its first EOS (whole seq if none)."""
|
||||
for i, tok in enumerate(seq):
|
||||
if tok in self.eos_ids:
|
||||
return seq[: i + 1]
|
||||
return list(seq)
|
||||
|
||||
def _diff_msg(self, prompt, ref, sgl):
|
||||
horizon = min(len(ref), len(sgl))
|
||||
first = next((j for j in range(horizon) if ref[j] != sgl[j]), horizon)
|
||||
@@ -246,10 +253,18 @@ class TestMlxReferenceCorrectness(CustomTestCase):
|
||||
for rid in rids:
|
||||
self.runner.remove_request(rid)
|
||||
|
||||
# Compare up to and including the first EOS. The horizon is fixed so
|
||||
# the batch composition never changes mid-run, which walks past EOS on
|
||||
# short answers -- and there the distribution is near-degenerate, so
|
||||
# batched and solo argmax can pick different tokens from a numerical
|
||||
# tie. That is float reduction order (a padded batched SDPA vs an
|
||||
# unpadded solo one), not state bleed: any cache crossover would show
|
||||
# up while the model still has an opinion. Measured on this fixture,
|
||||
# case 1 reaches EOS at index 2 and first differs at index 6.
|
||||
for i, (prompt, _, _) in enumerate(self.cases):
|
||||
self.assertEqual(
|
||||
batched[i], solo[i], self._diff_msg(prompt, solo[i], batched[i])
|
||||
)
|
||||
want = self._truncate_at_eos(solo[i])
|
||||
got = self._truncate_at_eos(batched[i])
|
||||
self.assertEqual(got, want, self._diff_msg(prompt, want, got))
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,728 @@
|
||||
"""Unit tests for MLX in-graph sampling (hardware_backend/mlx/sampling.py)."""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import unittest
|
||||
from collections import Counter
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=2, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=20, suite="stage-a-unit-test-mlx")
|
||||
|
||||
_HAS_MLX = importlib.util.find_spec("mlx") is not None
|
||||
_SKIP_REASON = "requires mlx"
|
||||
|
||||
if _HAS_MLX:
|
||||
import mlx.core as mx
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.sampling import (
|
||||
DEFAULT_SAMPLING_SEED,
|
||||
GREEDY_PARAMS,
|
||||
MAX_BOUNDED_TOP_K,
|
||||
MlxLogprobSpec,
|
||||
MlxSamplingParams,
|
||||
_candidate_width,
|
||||
_gumbel_noise,
|
||||
_murmur_hash32,
|
||||
all_greedy,
|
||||
compute_logprobs,
|
||||
sample_tokens,
|
||||
sanitize_logits,
|
||||
)
|
||||
|
||||
|
||||
def _reference_murmur3(seed: int, pos: int, col: int) -> int:
|
||||
"""Pure-Python MurmurHash3 mirroring the Triton kernel in
|
||||
sglang/kernels/ops/sampling/murmur_hash.py: blocks seed_low,
|
||||
seed_high, position, column; length-16 finalization; fmix32."""
|
||||
|
||||
def mix(h: int, k: int) -> int:
|
||||
k = (k * 0xCC9E2D51) & 0xFFFFFFFF
|
||||
k = ((k << 15) | (k >> 17)) & 0xFFFFFFFF
|
||||
k = (k * 0x1B873593) & 0xFFFFFFFF
|
||||
h ^= k
|
||||
h = ((h << 13) | (h >> 19)) & 0xFFFFFFFF
|
||||
return (h * 5 + 0xE6546B64) & 0xFFFFFFFF
|
||||
|
||||
seed &= 0xFFFFFFFFFFFFFFFF
|
||||
h = mix(0, seed & 0xFFFFFFFF)
|
||||
h = mix(h, (seed >> 32) & 0xFFFFFFFF)
|
||||
h = mix(h, pos & 0xFFFFFFFF)
|
||||
h = mix(h, col & 0xFFFFFFFF)
|
||||
h ^= 16
|
||||
h ^= h >> 16
|
||||
h = (h * 0x85EBCA6B) & 0xFFFFFFFF
|
||||
h ^= h >> 13
|
||||
h = (h * 0xC2B2AE35) & 0xFFFFFFFF
|
||||
h ^= h >> 16
|
||||
return h
|
||||
|
||||
|
||||
def _params(temperature=1.0, top_k=1 << 30, top_p=1.0, min_p=0.0, seed=None):
|
||||
return MlxSamplingParams(
|
||||
temperature=temperature, top_k=top_k, top_p=top_p, min_p=min_p, seed=seed
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestMurmurHashPort(CustomTestCase):
|
||||
def test_matches_pure_python_reference(self):
|
||||
"""Guards the mx uint32 port of the CUDA murmur kernel: any drift in
|
||||
wraparound/shift/block-order semantics changes seeded sampling."""
|
||||
seeds = [0, 1, 42, 2**31, 2**63 + 12345]
|
||||
positions = [0, 7, 1023, 2**31 - 1, 5]
|
||||
vocab = 64
|
||||
hashed = _murmur_hash32(seeds=seeds, positions=positions, vocab_size=vocab)
|
||||
mx.eval(hashed)
|
||||
for row, (seed, pos) in enumerate(zip(seeds, positions)):
|
||||
for col in (0, 1, vocab // 2, vocab - 1):
|
||||
self.assertEqual(
|
||||
int(hashed[row, col].item()),
|
||||
_reference_murmur3(seed, pos, col),
|
||||
msg=f"mismatch at seed={seed} pos={pos} col={col}",
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestSampleTokens(CustomTestCase):
|
||||
VOCAB = 32
|
||||
|
||||
def _logits(self, batch_size: int, key_int: int = 0) -> mx.array:
|
||||
return (
|
||||
mx.random.normal(shape=(batch_size, self.VOCAB), key=mx.random.key(key_int))
|
||||
* 3.0
|
||||
)
|
||||
|
||||
def _draw(self, logits, params, positions=None, n=200, key_start=100):
|
||||
"""Sample n times with distinct keys, return per-row token Counters."""
|
||||
batch_size = logits.shape[0]
|
||||
positions = positions if positions is not None else [5] * batch_size
|
||||
counters = [Counter() for _ in range(batch_size)]
|
||||
for i in range(n):
|
||||
toks = sample_tokens(
|
||||
last_logits=logits,
|
||||
params=params,
|
||||
positions=positions,
|
||||
key=mx.random.key(key_start + i),
|
||||
)
|
||||
mx.eval(toks)
|
||||
for row, t in enumerate(toks.tolist()):
|
||||
counters[row][int(t)] += 1
|
||||
return counters
|
||||
|
||||
@staticmethod
|
||||
def _reference_support(probs, top_k, top_p, min_p):
|
||||
"""Independent replica of the mask, in pure Python, on sorted probs."""
|
||||
order = sorted(range(len(probs)), key=lambda i: (-probs[i], i))
|
||||
keep, cum = [], 0.0
|
||||
for rank, idx in enumerate(order):
|
||||
p = probs[idx]
|
||||
masked = (
|
||||
rank >= min(top_k, len(probs))
|
||||
or cum > top_p
|
||||
or p < probs[order[0]] * min_p
|
||||
)
|
||||
cum += p
|
||||
if not masked:
|
||||
keep.append(idx)
|
||||
return set(keep)
|
||||
|
||||
def _probs(self, logits):
|
||||
probs = mx.softmax(logits.astype(mx.float32), axis=-1)
|
||||
mx.eval(probs)
|
||||
return probs[0].tolist()
|
||||
|
||||
def test_greedy_rows_match_argmax_in_mixed_batch(self):
|
||||
"""A greedy row must return exactly argmax even when other rows in
|
||||
the batch sample — guards the where() row-select and the sglang
|
||||
greedy convention (top_k == 1)."""
|
||||
logits = self._logits(3)
|
||||
expected = mx.argmax(logits, axis=-1).tolist()
|
||||
params = [_params(top_k=1), _params(temperature=0.7), _params(top_k=1)]
|
||||
counters = self._draw(logits, params, n=25)
|
||||
self.assertEqual(set(counters[0]), {expected[0]})
|
||||
self.assertEqual(set(counters[2]), {expected[2]})
|
||||
|
||||
def test_filter_supports_match_reference(self):
|
||||
"""The sampled support must stay inside the independently computed
|
||||
mask for each filter and for their combination — guards the rank
|
||||
mask, the nucleus exclusion, the min_p threshold, and the
|
||||
sorted->vocab index map."""
|
||||
cases = [
|
||||
("top_k=2", dict(top_k=2)),
|
||||
("top_p=0.6", dict(top_p=0.6)),
|
||||
("min_p=0.3", dict(min_p=0.3)),
|
||||
("top_k=8,top_p=0.7,min_p=0.05", dict(top_k=8, top_p=0.7, min_p=0.05)),
|
||||
]
|
||||
for label, kwargs in cases:
|
||||
with self.subTest(label):
|
||||
logits = self._logits(1, key_int=3)
|
||||
support = self._reference_support(
|
||||
self._probs(logits),
|
||||
kwargs.get("top_k", 1 << 30),
|
||||
kwargs.get("top_p", 1.0),
|
||||
kwargs.get("min_p", 0.0),
|
||||
)
|
||||
drawn = set(self._draw(logits, [_params(**kwargs)], n=400)[0])
|
||||
self.assertTrue(drawn <= support, f"{label}: extra {drawn - support}")
|
||||
self.assertTrue(drawn, f"{label}: nothing sampled")
|
||||
|
||||
def test_seeded_row_is_deterministic_and_key_independent(self):
|
||||
"""A row with sampling_seed must produce the same token regardless
|
||||
of the RNG key or batch composition — the murmur-gumbel path only
|
||||
depends on (seed, position, logits)."""
|
||||
logits = self._logits(2, key_int=5)
|
||||
seeded = _params(temperature=1.0, seed=1234)
|
||||
tok_solo = sample_tokens(
|
||||
last_logits=logits[:1],
|
||||
params=[seeded],
|
||||
positions=[9],
|
||||
key=mx.random.key(0),
|
||||
)
|
||||
tok_other_key = sample_tokens(
|
||||
last_logits=logits[:1],
|
||||
params=[seeded],
|
||||
positions=[9],
|
||||
key=mx.random.key(999),
|
||||
)
|
||||
tok_in_batch = sample_tokens(
|
||||
last_logits=logits,
|
||||
params=[seeded, _params(temperature=0.8)],
|
||||
positions=[9, 3],
|
||||
key=mx.random.key(7),
|
||||
)
|
||||
mx.eval(tok_solo, tok_other_key, tok_in_batch)
|
||||
self.assertEqual(tok_solo.tolist(), tok_other_key.tolist())
|
||||
self.assertEqual(int(tok_in_batch[0].item()), int(tok_solo[0].item()))
|
||||
|
||||
def test_seeded_row_unaffected_by_batchmate_filtering(self):
|
||||
"""A seeded row's token must not change when a batchmate triggers
|
||||
the top-k/top-p sort path — guards the vocab-id-space noise
|
||||
contract (regression: noise was applied in sorted-rank space when
|
||||
any row filtered, so batch composition changed seeded tokens)."""
|
||||
# Near-uniform seeded row: the Gumbel noise decides the token, so
|
||||
# a change of noise index space is guaranteed to show up.
|
||||
seeded_logits = (
|
||||
mx.random.normal(shape=(1, self.VOCAB), key=mx.random.key(8)) * 0.05
|
||||
)
|
||||
mate_logits = (
|
||||
mx.random.normal(shape=(1, self.VOCAB), key=mx.random.key(9)) * 3.0
|
||||
)
|
||||
logits = mx.concatenate([seeded_logits, mate_logits], axis=0)
|
||||
seeded = _params(seed=4321)
|
||||
solo = sample_tokens(
|
||||
last_logits=logits[:1],
|
||||
params=[seeded],
|
||||
positions=[6],
|
||||
key=mx.random.key(0),
|
||||
)
|
||||
with_filtering_mate = sample_tokens(
|
||||
last_logits=logits,
|
||||
params=[seeded, _params(temperature=1.2, top_k=2)],
|
||||
positions=[6, 11],
|
||||
key=mx.random.key(55),
|
||||
)
|
||||
mx.eval(solo, with_filtering_mate)
|
||||
self.assertEqual(int(with_filtering_mate[0].item()), int(solo[0].item()))
|
||||
|
||||
def test_bounded_top_k_picks_the_same_token_as_the_full_vocab_chain(self):
|
||||
"""The bounded top-K chain is an optimization, not a policy change:
|
||||
a seeded row must pick the same token whether or not the batch is
|
||||
eligible for it. A batchmate without a finite top_k pushes the
|
||||
whole batch back onto the full-vocab chain, so the same seeded row
|
||||
is sampled both ways here."""
|
||||
logits = self._logits(2, key_int=11)
|
||||
seeded = _params(top_k=4, seed=2024)
|
||||
bounded = sample_tokens(
|
||||
last_logits=logits[:1],
|
||||
params=[seeded],
|
||||
positions=[6],
|
||||
key=mx.random.key(0),
|
||||
)
|
||||
# min_p alone leaves top_k at TOP_K_ALL, so this batch falls back.
|
||||
full_vocab = sample_tokens(
|
||||
last_logits=logits,
|
||||
params=[seeded, _params(min_p=0.1)],
|
||||
positions=[6, 2],
|
||||
key=mx.random.key(3),
|
||||
)
|
||||
mx.eval(bounded, full_vocab)
|
||||
self.assertEqual(int(full_vocab[0].item()), int(bounded[0].item()))
|
||||
|
||||
def test_candidate_width_gates_the_bounded_chain(self):
|
||||
"""Only a batch whose widest top_k fits inside both the bound and
|
||||
the vocabulary may shrink the chain; anything else must return the
|
||||
full vocab size (which selects the scatter-back path)."""
|
||||
vocab = 4096
|
||||
for label, params, expected in [
|
||||
("under the bound", [_params(top_k=64)], 64),
|
||||
("at the bound", [_params(top_k=MAX_BOUNDED_TOP_K)], MAX_BOUNDED_TOP_K),
|
||||
("past the bound", [_params(top_k=MAX_BOUNDED_TOP_K + 1)], vocab),
|
||||
("no top_k (TOP_K_ALL)", [_params()], vocab),
|
||||
("top_k == vocab", [_params(top_k=vocab)], vocab),
|
||||
("widest row wins", [_params(top_k=4), _params(top_k=64)], 64),
|
||||
("one unbounded row", [_params(top_k=4), _params(top_p=0.9)], vocab),
|
||||
]:
|
||||
with self.subTest(label):
|
||||
self.assertEqual(_candidate_width(params, vocab), expected)
|
||||
|
||||
def test_seeded_row_varies_with_position(self):
|
||||
"""Positions feed the hash, so a fixed seed must not freeze the
|
||||
distribution across steps: over many positions the sampled tokens
|
||||
must not all collapse to one value (vocab of near-uniform probs)."""
|
||||
logits = mx.zeros((1, self.VOCAB)) # uniform distribution
|
||||
seeded = [_params(seed=77)]
|
||||
toks = set()
|
||||
for pos in range(40):
|
||||
t = sample_tokens(
|
||||
last_logits=logits,
|
||||
params=seeded,
|
||||
positions=[pos],
|
||||
key=mx.random.key(0),
|
||||
)
|
||||
mx.eval(t)
|
||||
toks.add(int(t[0].item()))
|
||||
self.assertGreater(len(toks), 5, toks)
|
||||
|
||||
def test_seeded_noise_is_finite(self):
|
||||
"""The uniform draw is clamped to [2**-32, 1 - 2**-24] before the
|
||||
double log: uint32(0xFFFFFFFF) rounds UP to 2**32 in float32, so an
|
||||
unclamped u can exceed 1 and make log(-log u) NaN — and u == 1 gives
|
||||
+inf, which would deterministically force that token."""
|
||||
# Sanity-check the hazard the clamp exists for.
|
||||
u_max = mx.array([0xFFFFFFFF], dtype=mx.uint32).astype(mx.float32) / float(
|
||||
0xFFFFFFFF
|
||||
)
|
||||
mx.eval(u_max)
|
||||
self.assertGreaterEqual(float(u_max.item()), 1.0)
|
||||
self.assertFalse(bool(mx.isfinite(-mx.log(-mx.log(u_max))).item()))
|
||||
|
||||
noise = _gumbel_noise(
|
||||
params=[_params(seed=1), _params(seed=2**63 - 1)],
|
||||
positions=[0, 4096],
|
||||
shape=(2, 1 << 16),
|
||||
key=mx.random.key(0),
|
||||
)
|
||||
mx.eval(noise)
|
||||
self.assertTrue(bool(mx.all(mx.isfinite(noise)).item()))
|
||||
|
||||
def test_seed_with_min_p_is_supported(self):
|
||||
"""seed + min_p is well defined here (the pytorch backend asserts on
|
||||
the combination): Gumbel-max over unnormalized masked weights is
|
||||
invariant to the missing renormalization, which is exactly the TODO
|
||||
at layers/sampler.py's multinomial_with_seed path."""
|
||||
logits = self._logits(1, key_int=4)
|
||||
support = self._reference_support(self._probs(logits), 1 << 30, 1.0, 0.3)
|
||||
tok = sample_tokens(
|
||||
last_logits=logits,
|
||||
params=[_params(seed=1234, min_p=0.3)],
|
||||
positions=[9],
|
||||
key=mx.random.key(0),
|
||||
)
|
||||
mx.eval(tok)
|
||||
self.assertIn(int(tok[0].item()), support)
|
||||
|
||||
def test_temperature_sharpens_distribution(self):
|
||||
"""Lower temperature must concentrate mass on the argmax token —
|
||||
guards the per-row temperature division (e.g. broadcasting bugs
|
||||
that apply one row's temperature to all rows)."""
|
||||
logits = self._logits(2, key_int=6)
|
||||
expected0 = int(mx.argmax(logits[0]).item())
|
||||
params = [_params(temperature=0.05), _params(temperature=5.0)]
|
||||
counters = self._draw(logits, params, n=200)
|
||||
self.assertGreater(counters[0][expected0] / 200.0, 0.95)
|
||||
self.assertGreater(len(counters[1]), 5, "high temp should spread mass")
|
||||
|
||||
def test_greedy_helpers(self):
|
||||
self.assertTrue(all_greedy([GREEDY_PARAMS, _params(top_k=1)]))
|
||||
self.assertFalse(all_greedy([GREEDY_PARAMS, _params(temperature=0.9)]))
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestSanitizeAndLogprobs(CustomTestCase):
|
||||
VOCAB = 16
|
||||
|
||||
def test_sanitize_matches_nan_to_num_semantics(self):
|
||||
"""Guards the port of sanitize_nan_logits' exact replacement values
|
||||
(+-1e30, not dtype extremes — temperature division would overflow
|
||||
dtype extremes back to inf and softmax them to NaN)."""
|
||||
import struct
|
||||
|
||||
def f32(v):
|
||||
return struct.unpack("f", struct.pack("f", v))[0]
|
||||
|
||||
x = mx.array([[1.0, float("nan"), float("inf"), -float("inf")]])
|
||||
out = sanitize_logits(x)
|
||||
mx.eval(out)
|
||||
self.assertEqual(out.tolist(), [[1.0, f32(-1e30), f32(1e30), f32(-1e30)]])
|
||||
|
||||
def test_logprobs_match_reference_and_row_shapes(self):
|
||||
"""compute_logprobs must equal log_softmax(logits/temp) per row and
|
||||
cut top-k / token-ids to each row's requested shape — guards the
|
||||
per-row temperature broadcast and the spec row alignment."""
|
||||
import math
|
||||
|
||||
logits = mx.random.normal(shape=(2, self.VOCAB), key=mx.random.key(11))
|
||||
params = [_params(temperature=0.5), _params(temperature=2.0)]
|
||||
tokens = mx.array([3, 7], dtype=mx.uint32)
|
||||
spec = MlxLogprobSpec(top_ks=(2, 0), token_ids=(None, (1, 4)))
|
||||
lp = compute_logprobs(logits, params, tokens, spec)
|
||||
mx.eval(*[a for a in [lp.chosen, lp.top_val, lp.top_idx] if a is not None])
|
||||
|
||||
raw = logits.tolist()
|
||||
for row, temp in ((0, 0.5), (1, 2.0)):
|
||||
scaled = [v / temp for v in raw[row]]
|
||||
m = max(scaled)
|
||||
lse = m + math.log(sum(math.exp(v - m) for v in scaled))
|
||||
ref = [v - lse for v in scaled]
|
||||
chosen_token = int(tokens[row].item())
|
||||
self.assertAlmostEqual(
|
||||
float(lp.chosen[row].item()), ref[chosen_token], places=4
|
||||
)
|
||||
if row == 0:
|
||||
expect_top = sorted(ref, reverse=True)[:2]
|
||||
got = lp.top_val[row].tolist()[:2]
|
||||
for a, b in zip(got, expect_top):
|
||||
self.assertAlmostEqual(a, b, places=4)
|
||||
if row == 1:
|
||||
got = lp.token_ids_val[1].tolist()
|
||||
self.assertAlmostEqual(got[0], ref[1], places=4)
|
||||
self.assertAlmostEqual(got[1], ref[4], places=4)
|
||||
self.assertIsNone(lp.token_ids_val[0])
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestRunnerSelectTokens(CustomTestCase):
|
||||
"""_select_tokens_with_logprobs lifecycle on a bare runner (object.__new__)."""
|
||||
|
||||
class _FakeCache:
|
||||
def __init__(self, offset):
|
||||
self.offset = offset
|
||||
|
||||
class _FakeLayout:
|
||||
has_auxiliary_state = False
|
||||
first_attention_layer_index = 0
|
||||
|
||||
def _runner(self, enable_sampling):
|
||||
from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
|
||||
|
||||
runner = object.__new__(MlxModelRunner)
|
||||
runner._enable_sampling = enable_sampling
|
||||
runner._cache_layout = self._FakeLayout()
|
||||
runner._req_sampling = {}
|
||||
runner._rng_key = mx.random.key(0) if enable_sampling else None
|
||||
return runner
|
||||
|
||||
def test_disabled_and_greedy_paths_consume_no_rng(self):
|
||||
"""Flag-off and all-greedy batches must return exact argmax and
|
||||
leave the RNG key untouched — guards the byte-exact greedy
|
||||
contract that the e2e temp=0 test relies on."""
|
||||
logits = mx.random.normal(shape=(2, 16), key=mx.random.key(1))
|
||||
expected = mx.argmax(logits, axis=-1).tolist()
|
||||
caches = [[self._FakeCache(4)], [self._FakeCache(9)]]
|
||||
|
||||
disabled = self._runner(enable_sampling=False)
|
||||
toks = disabled._select_tokens_with_logprobs(logits, ["a", "b"], caches)[0]
|
||||
self.assertEqual(toks.tolist(), expected)
|
||||
|
||||
enabled = self._runner(enable_sampling=True)
|
||||
enabled._req_sampling = {"a": GREEDY_PARAMS, "b": _params(top_k=1)}
|
||||
key_before = enabled._rng_key
|
||||
toks = enabled._select_tokens_with_logprobs(logits, ["a", "b"], caches)[0]
|
||||
self.assertEqual(toks.tolist(), expected)
|
||||
self.assertIs(enabled._rng_key, key_before)
|
||||
|
||||
def test_discarded_chunk_without_trunk_stays_greedy_and_consumes_no_rng(self):
|
||||
"""A needs_logits=False chunk on a model without a headless trunk
|
||||
must not sample: consuming RNG for a discarded token would make
|
||||
the final output depend on prefill chunking."""
|
||||
|
||||
def full_model_only(input_ids, cache=None):
|
||||
return mx.zeros((1, input_ids.shape[1], 16))
|
||||
|
||||
runner = self._runner(enable_sampling=True)
|
||||
runner.model = full_model_only # no .model attr -> no trunk
|
||||
runner._req_sampling = {"a": _params(temperature=1.0)}
|
||||
key_before = runner._rng_key
|
||||
tok, lazy_logprobs = runner._forward_lazy_token(
|
||||
mx.array([[3, 4]], dtype=mx.int32),
|
||||
[self._FakeCache(2)],
|
||||
needs_logits=False,
|
||||
req_id="a",
|
||||
)
|
||||
self.assertIsNone(lazy_logprobs)
|
||||
mx.eval(tok)
|
||||
self.assertEqual(tok.tolist(), [0]) # argmax of zeros
|
||||
self.assertIs(runner._rng_key, key_before)
|
||||
|
||||
def test_logit_edits_gate_greedy_and_sampled_and_logprobs(self):
|
||||
"""An additive -inf edit row must exclude a token from greedy argmax,
|
||||
from sampling, AND from the reported logprob distribution — guards
|
||||
the edits-before-selection ordering (a regression that samples raw
|
||||
logits would pass every other test on near-uniform inputs)."""
|
||||
runner = self._runner(enable_sampling=True)
|
||||
runner._req_sampling = {"g": GREEDY_PARAMS, "s": _params(temperature=1.0)}
|
||||
caches = [[self._FakeCache(4)], [self._FakeCache(4)]]
|
||||
logits = mx.zeros((2, 8))
|
||||
logits = mx.put_along_axis(
|
||||
logits,
|
||||
mx.array([[7], [7]], dtype=mx.uint32),
|
||||
mx.array([[5.0], [5.0]]),
|
||||
axis=-1,
|
||||
) # token 7 dominates both rows
|
||||
edits = mx.zeros((2, 8))
|
||||
edits = mx.put_along_axis(
|
||||
edits,
|
||||
mx.array([[7], [7]], dtype=mx.uint32),
|
||||
mx.array([[-float("inf")], [-float("inf")]]),
|
||||
axis=-1,
|
||||
) # ...but is masked out for both
|
||||
spec = MlxLogprobSpec(top_ks=(1, 1), token_ids=(None, None))
|
||||
for _ in range(10):
|
||||
tokens, lp = runner._select_tokens_with_logprobs(
|
||||
logits, ["g", "s"], caches, edits, spec
|
||||
)
|
||||
mx.eval(tokens, lp.chosen, lp.top_val)
|
||||
self.assertNotIn(7, tokens.tolist())
|
||||
self.assertNotIn(7, [row[0] for row in lp.top_idx.tolist()])
|
||||
|
||||
def test_seed_is_gated_on_deterministic_inference(self):
|
||||
"""Upstream seed contract: SamplingBatchInfo only populates
|
||||
sampling_seed under --enable-deterministic-inference, and then seeds
|
||||
every row (default 42). A per-request seed outside that flag is
|
||||
ignored by every other backend, so it is ignored here too."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
def make_req(sampling_seed):
|
||||
return SimpleNamespace(
|
||||
sampling_params=SimpleNamespace(
|
||||
temperature=0.8,
|
||||
top_k=1 << 30,
|
||||
top_p=1.0,
|
||||
min_p=0.0,
|
||||
sampling_seed=sampling_seed,
|
||||
frequency_penalty=0.0,
|
||||
presence_penalty=0.0,
|
||||
repetition_penalty=1.0,
|
||||
)
|
||||
)
|
||||
|
||||
self.assertIsNone(MlxSamplingParams.from_req(make_req(7)).seed)
|
||||
self.assertIsNone(
|
||||
MlxSamplingParams.from_req(make_req(None), deterministic_seeding=False).seed
|
||||
)
|
||||
self.assertEqual(
|
||||
MlxSamplingParams.from_req(make_req(7), deterministic_seeding=True).seed, 7
|
||||
)
|
||||
self.assertEqual(
|
||||
MlxSamplingParams.from_req(make_req(None), deterministic_seeding=True).seed,
|
||||
DEFAULT_SAMPLING_SEED,
|
||||
)
|
||||
|
||||
def test_chained_decode_keeps_logit_bias(self):
|
||||
"""A chained decode step must keep applying the batch's static
|
||||
logit_bias rows — regression: the chained path passed edits=None,
|
||||
silently dropping the bias after the first (fresh) step."""
|
||||
runner = self._runner(enable_sampling=True)
|
||||
runner._req_sampling = {"a": GREEDY_PARAMS}
|
||||
runner._req_caches = {"a": [self._FakeCache(3)]}
|
||||
runner._req_token_ids = {"a": [1]}
|
||||
# token 2 dominates; the edit row bans it -> argmax must fall to 1
|
||||
logits = mx.array([[0.0, 3.0, 5.0, 0.0]])
|
||||
runner._decode_with_batched_attention = lambda caches, x, rids: logits
|
||||
edits = mx.array([[0.0, 0.0, -float("inf"), 0.0]])
|
||||
|
||||
fresh = runner.decode_batch_start(["a"], edit_rows=edits)
|
||||
chained = runner.decode_batch_start_chained(fresh)
|
||||
mx.eval(fresh.lazy_tokens, chained.lazy_tokens)
|
||||
self.assertEqual(fresh.lazy_tokens.tolist(), [1])
|
||||
self.assertEqual(chained.lazy_tokens.tolist(), [1])
|
||||
|
||||
def test_logits_hook_bridge_roundtrip(self):
|
||||
"""The custom-logit-processor hook must see materialized float32
|
||||
logits and its in-place edits must re-enter the graph — guards the
|
||||
mx->numpy->mx bridge (a copy-semantics change would drop edits)."""
|
||||
runner = self._runner(enable_sampling=True)
|
||||
logits = mx.zeros((1, 8), dtype=mx.bfloat16)
|
||||
|
||||
def hook(arr):
|
||||
assert arr.dtype.name == "float32"
|
||||
arr[0, 5] = 99.0
|
||||
return arr
|
||||
|
||||
edited = runner._run_logits_hook(logits, hook)
|
||||
mx.eval(edited)
|
||||
self.assertEqual(int(mx.argmax(edited, axis=-1)[0].item()), 5)
|
||||
|
||||
def test_sampling_path_advances_rng_key(self):
|
||||
"""Consecutive sampling builds must consume distinct keys, or every
|
||||
chained decode step would draw identical noise."""
|
||||
logits = mx.zeros((1, 16))
|
||||
runner = self._runner(enable_sampling=True)
|
||||
runner._req_sampling = {"a": _params(temperature=1.0)}
|
||||
caches = [[self._FakeCache(3)]]
|
||||
toks = set()
|
||||
for _ in range(20):
|
||||
t = runner._select_tokens_with_logprobs(logits, ["a"], caches)[0]
|
||||
mx.eval(t)
|
||||
toks.add(int(t[0].item()))
|
||||
self.assertGreater(len(toks), 3, toks)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestWorkerSamplingExtras(CustomTestCase):
|
||||
"""Worker-side builders: logit-edit rows, logprob specs, output assembly."""
|
||||
|
||||
VOCAB = 8
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
from sglang.srt.runtime_context import get_context
|
||||
|
||||
# The worker reads --mlx-enable-sampling off the device config bag,
|
||||
# which fails closed before a publish.
|
||||
cls._config = get_context().override_server_args(mlx_enable_sampling=True)
|
||||
cls._config.install()
|
||||
cls.addClassCleanup(cls._config.restore)
|
||||
|
||||
@staticmethod
|
||||
def _worker():
|
||||
from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker
|
||||
|
||||
return MlxTpModelWorker.__new__(MlxTpModelWorker)
|
||||
|
||||
def _batch(self, sinfo, n=2, return_logprob=False, has_grammar=False):
|
||||
from types import SimpleNamespace
|
||||
|
||||
return SimpleNamespace(
|
||||
reqs=[
|
||||
SimpleNamespace(
|
||||
rid=f"r{i}", return_logprob=return_logprob, grammar=None
|
||||
)
|
||||
for i in range(n)
|
||||
],
|
||||
sampling_info=sinfo,
|
||||
return_logprob=return_logprob,
|
||||
top_logprobs_nums=None,
|
||||
token_ids_logprobs=None,
|
||||
has_grammar=has_grammar,
|
||||
)
|
||||
|
||||
def test_edit_rows_combine_grammar_mask_and_bias(self):
|
||||
"""The grammar mask must be applied through the backend's own
|
||||
apply_vocab_mask on a zeros base and summed with logit_bias —
|
||||
guards the backend-agnostic zeros-trick, the combine order, and the
|
||||
ForwardBatch.init_new grammars-population mirror (regression: the
|
||||
MLX paths never build a ForwardBatch, so sinfo.grammars stayed None
|
||||
and live grammar objects produced no mask at all)."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
import torch
|
||||
|
||||
class FakeGrammar:
|
||||
def apply_vocab_mask(self, logits, vocab_mask):
|
||||
logits[0, 3] = -float("inf") # row 0 forbids token 3
|
||||
|
||||
sinfo = SimpleNamespace(
|
||||
grammars=None, # not yet populated, as on the real MLX path
|
||||
logit_bias=torch.zeros(2, self.VOCAB).index_put_(
|
||||
(torch.tensor([1]), torch.tensor([5])), torch.tensor([2.5])
|
||||
),
|
||||
vocab_size=self.VOCAB,
|
||||
grammar_mask=None,
|
||||
)
|
||||
|
||||
def update_mask():
|
||||
sinfo.grammar_mask = SimpleNamespace(
|
||||
grammar=FakeGrammar(), vocab_mask=torch.zeros(2, 1)
|
||||
)
|
||||
|
||||
sinfo.update_regex_vocab_mask = update_mask
|
||||
|
||||
batch = self._batch(sinfo, has_grammar=True)
|
||||
batch.reqs[0].grammar = object()
|
||||
rows = self._worker()._build_logit_edit_rows(batch)
|
||||
self.assertEqual(
|
||||
[g is not None for g in sinfo.grammars],
|
||||
[True, False],
|
||||
"worker must mirror ForwardBatch.init_new's grammars population",
|
||||
)
|
||||
mx.eval(rows["r0"], rows["r1"])
|
||||
self.assertEqual(rows["r0"].tolist()[3], -float("inf"))
|
||||
self.assertEqual(rows["r1"].tolist()[5], 2.5)
|
||||
self.assertIsNone(sinfo.grammar_mask, "mask must be released after use")
|
||||
|
||||
def test_edit_rows_none_when_nothing_to_edit(self):
|
||||
from types import SimpleNamespace
|
||||
|
||||
sinfo = SimpleNamespace(grammars=None, logit_bias=None, vocab_size=self.VOCAB)
|
||||
self.assertIsNone(self._worker()._build_logit_edit_rows(self._batch(sinfo)))
|
||||
|
||||
def test_logprob_spec_subset_alignment(self):
|
||||
"""Spec rows must align to the rid subset order, not batch order —
|
||||
guards mixed-batch decode sub-batches."""
|
||||
from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker
|
||||
|
||||
rows = {"a": (3, None), "c": (0, (7, 9))}
|
||||
spec = MlxTpModelWorker._logprob_spec_for(rows, ["c", "b", "a"])
|
||||
self.assertEqual(spec.top_ks, (0, 0, 3))
|
||||
self.assertEqual(spec.token_ids, ((7, 9), None, None))
|
||||
self.assertIsNone(MlxTpModelWorker._logprob_spec_for(rows, ["x"]))
|
||||
|
||||
@unittest.skipUnless(
|
||||
importlib.util.find_spec("xgrammar") is not None, "requires xgrammar"
|
||||
)
|
||||
def test_xgrammar_wrapper_supports_cpu_logits(self):
|
||||
"""The MLX edit-row builder feeds CPU logits to the grammar
|
||||
backend's apply_vocab_mask — regression: the xgrammar wrapper
|
||||
raised 'Unsupported device: cpu' (its dispatch stopped at
|
||||
cuda/xpu/musa/npu), so every grammar request crashed the worker."""
|
||||
import math
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
|
||||
from sglang.srt.constrained.xgrammar_backend import XGrammarGrammar
|
||||
|
||||
logits = torch.zeros(1, 40)
|
||||
blocks = math.ceil(40 / 32)
|
||||
bitmask = torch.full((1, blocks), -1, dtype=torch.int32)
|
||||
bitmask[0, 0] = int(np.int32(np.uint32(0xFFFFFFFF & ~(1 << 7))))
|
||||
XGrammarGrammar.apply_vocab_mask(None, logits, bitmask)
|
||||
self.assertEqual(logits[0, 7].item(), -float("inf"))
|
||||
self.assertEqual(logits[0, 6].item(), 0.0)
|
||||
|
||||
def test_assemble_logprob_output_matches_scheduler_contract(self):
|
||||
"""Field shapes must survive the scheduler's move_logprobs_to_cpu
|
||||
(`.tolist()` on the batch tensor and on every per-row val/idx entry)
|
||||
and add_logprob_return_values indexing — guards the external
|
||||
LogitsProcessorOutput consumption contract, including rows without
|
||||
logprob requests getting empty-but-tolistable fills."""
|
||||
from types import SimpleNamespace
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker
|
||||
|
||||
step_rows = {"a": (-1.5, [-0.1, -0.2], [4, 2], [-3.0], [9])}
|
||||
reqs = [SimpleNamespace(rid="a"), SimpleNamespace(rid="b")]
|
||||
out = MlxTpModelWorker._assemble_logprob_output(step_rows, reqs)
|
||||
|
||||
self.assertEqual(out.next_token_logprobs.tolist(), [-1.5, 0.0])
|
||||
self.assertEqual(
|
||||
[v.tolist() for v in out.next_token_top_logprobs_val],
|
||||
[[-0.10000000149011612, -0.20000000298023224], []],
|
||||
)
|
||||
self.assertEqual(
|
||||
[v.tolist() for v in out.next_token_top_logprobs_idx], [[4, 2], []]
|
||||
)
|
||||
self.assertEqual(
|
||||
[v.tolist() for v in out.next_token_token_ids_logprobs_val],
|
||||
[[-3.0], []],
|
||||
)
|
||||
self.assertEqual(out.next_token_token_ids_logprobs_idx, [[9], []])
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -129,6 +129,7 @@ class TestOverlapLoopStampsLaunchTs(unittest.TestCase):
|
||||
from sglang.srt.hardware_backend.mlx.scheduler_mixin import (
|
||||
SchedulerMlxOverlapMixin,
|
||||
)
|
||||
from sglang.srt.hardware_backend.mlx.tp_worker import MlxLaunch
|
||||
|
||||
scheduler = self._make_scheduler(recv_side_effect=[[], _StopLoop()])
|
||||
|
||||
@@ -148,7 +149,13 @@ class TestOverlapLoopStampsLaunchTs(unittest.TestCase):
|
||||
scheduler.tp_worker.async_forward_batch_generation_mlx.side_effect = (
|
||||
lambda _batch: (
|
||||
events.append("forward"),
|
||||
(None, [], [], None, "extend"),
|
||||
MlxLaunch(
|
||||
lazy_tokens=None,
|
||||
prefills=[],
|
||||
extends=[],
|
||||
decode=None,
|
||||
mode="extend",
|
||||
),
|
||||
)[1]
|
||||
)
|
||||
|
||||
@@ -173,6 +180,7 @@ class TestOverlapLoopStampsLaunchTs(unittest.TestCase):
|
||||
from sglang.srt.hardware_backend.mlx.scheduler_mixin import (
|
||||
SchedulerMlxOverlapMixin,
|
||||
)
|
||||
from sglang.srt.hardware_backend.mlx.tp_worker import MlxLaunch
|
||||
|
||||
# Iteration 1: fresh decode launch. Iteration 2: chain a second
|
||||
# decode on top of it. Iteration 3: stop.
|
||||
@@ -193,16 +201,22 @@ class TestOverlapLoopStampsLaunchTs(unittest.TestCase):
|
||||
scheduler.get_next_batch_to_run.return_value = plan
|
||||
|
||||
pending_decode = MagicMock()
|
||||
scheduler.tp_worker.async_forward_batch_generation_mlx.return_value = (
|
||||
MagicMock(),
|
||||
[],
|
||||
[],
|
||||
pending_decode,
|
||||
"decode",
|
||||
scheduler.tp_worker.async_forward_batch_generation_mlx.return_value = MlxLaunch(
|
||||
lazy_tokens=MagicMock(),
|
||||
prefills=[],
|
||||
extends=[],
|
||||
decode=pending_decode,
|
||||
mode="decode",
|
||||
)
|
||||
scheduler.tp_worker.async_chained_decode_mlx.side_effect = lambda _decode: (
|
||||
events.append("chained_forward"),
|
||||
(MagicMock(), [], [], MagicMock(), "decode"),
|
||||
MlxLaunch(
|
||||
lazy_tokens=MagicMock(),
|
||||
prefills=[],
|
||||
extends=[],
|
||||
decode=MagicMock(),
|
||||
mode="decode",
|
||||
),
|
||||
)[1]
|
||||
|
||||
launch_times = iter((1.0, 2.0))
|
||||
|
||||
@@ -7,10 +7,13 @@ pin the three seams that make such models work on the MLX backend:
|
||||
1. The attention contract accepts ``sm_scale`` and exposes per-layer window
|
||||
sizes read from the mlx-lm container convention (``layer_types`` +
|
||||
``window_size``).
|
||||
2. The cache shims' ``make_mask`` mirrors mlx_lm's
|
||||
``cache.create_attention_mask`` exactly — in particular ``window_size``
|
||||
2. The cache shims' ``make_mask`` matches mlx_lm's
|
||||
``cache.create_attention_mask`` semantically — in particular ``window_size``
|
||||
must produce a banded mask (including for N == 1) instead of being
|
||||
silently dropped, or sliding-window layers degrade to full attention.
|
||||
Where the window provably cannot bind (``offset + N <= window_size``) the
|
||||
band equals plain causal, and the shims return the cheap form instead, as
|
||||
mlx_lm's own ``RotatingKVCache.make_mask`` does.
|
||||
3. ``MLXAttentionWrapper._batched_decode`` applies the window by truncating
|
||||
each request's KV to the trailing window, passes ``sinks`` through, and
|
||||
uses the contract scale helper.
|
||||
@@ -173,7 +176,22 @@ class TestGptOssAttentionContract(CustomTestCase):
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestShimMakeMask(CustomTestCase):
|
||||
"""The shims must return exactly what mlx_lm's own KVCache.make_mask returns."""
|
||||
"""The shims must be semantically equal to mlx_lm's own KVCache.make_mask.
|
||||
|
||||
Equal *content*, not equal representation: where the window provably
|
||||
cannot bind the shims return the cheap ``"causal"`` / ``None`` form that
|
||||
mlx_lm's RotatingKVCache.make_mask also returns, so the comparison
|
||||
densifies both sides first.
|
||||
"""
|
||||
|
||||
def _dense(self, mask, N, offset):
|
||||
"""Dense boolean form of any of the three mask representations."""
|
||||
if mask is None:
|
||||
return mx.ones((N, offset + N), dtype=mx.bool_)
|
||||
if isinstance(mask, str):
|
||||
self.assertEqual(mask, "causal")
|
||||
return create_causal_mask(N, offset)
|
||||
return mask
|
||||
|
||||
def _shims(self, offset):
|
||||
contig = ContiguousAttentionKVCache(
|
||||
@@ -185,14 +203,13 @@ class TestShimMakeMask(CustomTestCase):
|
||||
)
|
||||
return (AttentionOffsetCache(offset=offset), contig, pool_backed)
|
||||
|
||||
def _assert_same_mask(self, got, ref, msg):
|
||||
if ref is None or isinstance(ref, str):
|
||||
self.assertEqual(got, ref, msg)
|
||||
else:
|
||||
self.assertTrue(
|
||||
isinstance(got, mx.array) and mx.array_equal(got, ref).item(),
|
||||
msg,
|
||||
)
|
||||
def _assert_same_mask(self, got, ref, msg, N, offset):
|
||||
self.assertTrue(
|
||||
mx.array_equal(
|
||||
self._dense(got, N, offset), self._dense(ref, N, offset)
|
||||
).item(),
|
||||
msg,
|
||||
)
|
||||
|
||||
def test_shims_match_mlx_lm_reference(self):
|
||||
cases = [
|
||||
@@ -213,8 +230,34 @@ class TestShimMakeMask(CustomTestCase):
|
||||
ref,
|
||||
f"{type(shim).__name__} mismatch for N={N} offset={offset} "
|
||||
f"window={window} return_array={return_array}",
|
||||
N,
|
||||
offset,
|
||||
)
|
||||
|
||||
def test_non_binding_window_returns_the_cheap_mask(self):
|
||||
# offset + N <= window: no query can reach past the window, so the
|
||||
# band equals plain causal and materialising it only costs time (a
|
||||
# mask array forces sdpa off its fused causal path, ~2x per layer).
|
||||
self.assertIsNone(make_attention_mask(1, 0, window_size=4))
|
||||
self.assertIsNone(make_attention_mask(1, 3, window_size=4))
|
||||
self.assertEqual(make_attention_mask(4, 0, window_size=4), "causal")
|
||||
# ...and one position past the boundary the band is required again.
|
||||
self.assertIsInstance(make_attention_mask(4, 1, window_size=4), mx.array)
|
||||
self.assertIsInstance(make_attention_mask(1, 4, window_size=4), mx.array)
|
||||
|
||||
def test_non_binding_window_matches_the_band_it_replaces(self):
|
||||
# The shortcut is only legal because the two forms are elementwise
|
||||
# identical; pin that against mlx_lm's own band builder.
|
||||
for N, offset, window in ((1, 0, 4), (1, 3, 4), (4, 0, 4), (8, 0, 16)):
|
||||
band = create_causal_mask(N, offset, window_size=window)
|
||||
cheap = self._dense(
|
||||
make_attention_mask(N, offset, window_size=window), N, offset
|
||||
)
|
||||
self.assertTrue(
|
||||
mx.array_equal(band, cheap).item(),
|
||||
f"N={N} offset={offset} window={window}",
|
||||
)
|
||||
|
||||
def test_windowed_mask_is_banded_including_self(self):
|
||||
# Query at absolute position 6 with W=4 may attend to keys 3..6
|
||||
# (j in [i - W + 1, i], the window includes the query itself).
|
||||
|
||||
@@ -0,0 +1,380 @@
|
||||
"""Unit tests for sliding-window layers on the MLX radix/pool KV path.
|
||||
|
||||
The shared ``MlxAttentionKVPool`` stores full-attention layers only,
|
||||
sliding-window layers keep window-bounded per-request storage, and a
|
||||
radix prefix hit on an SWA model recomputes the whole prefix. Scheduler
|
||||
bookkeeping stays in the unclamped coordinates, so these tests drive
|
||||
``MlxModelRunner`` directly with hand-built slot ids, mirroring how the
|
||||
tp_worker calls it.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=10, suite="stage-a-unit-test-mlx")
|
||||
|
||||
_HAS_MLX = (
|
||||
importlib.util.find_spec("mlx") is not None
|
||||
and importlib.util.find_spec("mlx_lm") is not None
|
||||
)
|
||||
_SKIP_REASON = "requires mlx + mlx_lm"
|
||||
|
||||
if _HAS_MLX:
|
||||
import mlx.core as mx
|
||||
from mlx_lm.models import gpt_oss
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.aot import (
|
||||
MlxAOTKernelContext,
|
||||
MlxAOTKernelSet,
|
||||
MlxAOTRoPEContext,
|
||||
MlxAOTRoPEKernel,
|
||||
)
|
||||
from sglang.srt.hardware_backend.mlx.kv_cache import (
|
||||
BatchedDecodeContext,
|
||||
ContiguousAttentionKVCache,
|
||||
MlxAttentionKVPool,
|
||||
MLXAttentionWrapper,
|
||||
WindowedAttentionKVCache,
|
||||
find_attention_layers,
|
||||
get_layer_window_sizes,
|
||||
patch_model_attention,
|
||||
)
|
||||
from sglang.srt.hardware_backend.mlx.kv_cache.layout import MlxModelCacheLayout
|
||||
from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
|
||||
|
||||
TINY_WINDOW = 8
|
||||
|
||||
|
||||
def _tiny_gpt_oss_model():
|
||||
"""Randomly initialized 4-layer gpt_oss with alternating sliding/full layers.
|
||||
|
||||
Mirrors test_windowed_kv_cache.py's builder (kept local: the registered
|
||||
unit-test directory is not an importable package).
|
||||
"""
|
||||
args = gpt_oss.ModelArgs(
|
||||
num_hidden_layers=4,
|
||||
num_local_experts=8,
|
||||
num_experts_per_tok=2,
|
||||
vocab_size=128,
|
||||
hidden_size=64,
|
||||
intermediate_size=64,
|
||||
head_dim=16,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=2,
|
||||
sliding_window=TINY_WINDOW,
|
||||
rope_theta=150000,
|
||||
rope_scaling={
|
||||
"rope_type": "yarn",
|
||||
"factor": 32.0,
|
||||
"beta_fast": 32.0,
|
||||
"beta_slow": 1.0,
|
||||
"original_max_position_embeddings": 4096,
|
||||
"truncate": False,
|
||||
},
|
||||
)
|
||||
return gpt_oss.Model(args)
|
||||
|
||||
|
||||
def _stub_runner(model, disable_radix_cache, pool_size=64):
|
||||
"""Surgically build a runner around an already-loaded tiny model."""
|
||||
layers, attrs = find_attention_layers(model)
|
||||
runner = MlxModelRunner.__new__(MlxModelRunner)
|
||||
runner.model = model
|
||||
runner.disable_radix_cache = disable_radix_cache
|
||||
runner._cache_layout = MlxModelCacheLayout.from_attention_discovery(
|
||||
layers, attrs, layer_window_sizes=get_layer_window_sizes(model)
|
||||
)
|
||||
runner._max_seq_len = 64
|
||||
runner._cache_pool = []
|
||||
runner._req_caches = {}
|
||||
runner._req_token_ids = {}
|
||||
runner._req_sampling = {}
|
||||
runner._req_pool_idx = {}
|
||||
runner._req_synced_offset = {}
|
||||
runner._req_to_token_pool = None
|
||||
runner._attention_kv_pool = None
|
||||
runner._decode_step_ct = 0
|
||||
runner._clear_steps = 0
|
||||
runner._aot_kernels = MlxAOTKernelSet()
|
||||
runner._pool_size = pool_size
|
||||
if not disable_radix_cache:
|
||||
runner.init_cache_pools(None)
|
||||
return runner
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestSwaLayoutAndPoolContract(CustomTestCase):
|
||||
def _layout(self, with_windows=True):
|
||||
model = _tiny_gpt_oss_model()
|
||||
layers, attrs = find_attention_layers(model)
|
||||
return MlxModelCacheLayout.from_attention_discovery(
|
||||
layers, attrs, get_layer_window_sizes(model) if with_windows else None
|
||||
)
|
||||
|
||||
def test_partition_and_dense_full_pool_index(self):
|
||||
layout = self._layout()
|
||||
self.assertEqual(layout.attention_layer_indices, (0, 1, 2, 3))
|
||||
self.assertEqual(layout.swa_attention_layer_indices, (0, 2))
|
||||
self.assertEqual(layout.full_attention_layer_indices, (1, 3))
|
||||
# Dense over full layers only, so it differs from the cache index.
|
||||
self.assertEqual(layout.full_kv_pool_index_by_layer, {1: 0, 3: 1})
|
||||
self.assertEqual(layout.attention_pool_index_by_layer, {0: 0, 1: 1, 2: 2, 3: 3})
|
||||
with self.assertRaises(KeyError):
|
||||
layout.full_kv_pool_index(0)
|
||||
|
||||
# Without a window map the two indices coincide (pre-SWA behavior).
|
||||
plain = self._layout(with_windows=False)
|
||||
self.assertFalse(plain.has_sliding_window_layers)
|
||||
self.assertEqual(
|
||||
plain.full_kv_pool_index_by_layer, plain.attention_pool_index_by_layer
|
||||
)
|
||||
|
||||
def test_sliding_window_model_gets_no_pool(self):
|
||||
# An SWA prefix hit recomputes the prefix instead of gathering it, so
|
||||
# the shared pool would have no reader. Allocating it would burn the
|
||||
# whole auto-sized KV budget on a write-only buffer.
|
||||
runner = _stub_runner(_tiny_gpt_oss_model(), disable_radix_cache=False)
|
||||
self.assertTrue(runner._cache_layout.has_sliding_window_layers)
|
||||
self.assertIsNone(runner._attention_kv_pool)
|
||||
# The layer-type split still resolves -- it is the seam the shared
|
||||
# window-aware SWA pool will build on.
|
||||
self.assertEqual(runner._cache_layout.full_kv_pool_index_by_layer, {1: 0, 3: 1})
|
||||
|
||||
def test_pool_covers_every_layer_without_sliding_windows(self):
|
||||
model = _tiny_gpt_oss_model()
|
||||
layers, attrs = find_attention_layers(model)
|
||||
runner = _stub_runner(model, disable_radix_cache=True)
|
||||
runner._cache_layout = MlxModelCacheLayout.from_attention_discovery(
|
||||
layers, attrs
|
||||
)
|
||||
runner.disable_radix_cache = False
|
||||
runner.init_cache_pools(None)
|
||||
self.assertEqual(runner._attention_kv_pool.num_layers, 4)
|
||||
self.assertEqual(runner._attention_kv_pool.pool_size, 65)
|
||||
|
||||
def test_all_sliding_window_model_gets_no_pool(self):
|
||||
# An all-SWA model has nothing to pool. Pool construction must skip
|
||||
# out, and pool sizing must still land on a finite slot count rather
|
||||
# than dividing by zero bytes per slot.
|
||||
runner = _stub_runner(_tiny_gpt_oss_model(), disable_radix_cache=False)
|
||||
layers, attrs = find_attention_layers(runner.model)
|
||||
runner._cache_layout = MlxModelCacheLayout.from_attention_discovery(
|
||||
layers, attrs, {idx: TINY_WINDOW for idx in range(4)}
|
||||
)
|
||||
self.assertEqual(runner._cache_layout.full_attention_layer_indices, ())
|
||||
self.assertEqual(runner._cache_layout.full_kv_pool_index_by_layer, {})
|
||||
|
||||
runner._attention_kv_pool = None
|
||||
runner.init_cache_pools(None)
|
||||
self.assertIsNone(runner._attention_kv_pool)
|
||||
|
||||
runner._mem_fraction_static = 0.5
|
||||
self.assertGreater(runner._compute_pool_size(None), 0)
|
||||
|
||||
def test_sliding_flag_without_window_map_still_rejected(self):
|
||||
model = _tiny_gpt_oss_model()
|
||||
patch_model_attention(model)
|
||||
model.model.layers[0].self_attn._inner.is_sliding = True
|
||||
runner = _stub_runner(model, disable_radix_cache=True)
|
||||
# With the container window map the flagged layer is bounded: fine.
|
||||
runner._get_attn_config()
|
||||
# Without a resolvable window the layer cannot be bounded: reject.
|
||||
layers, attrs = find_attention_layers(model)
|
||||
runner._cache_layout = MlxModelCacheLayout.from_attention_discovery(
|
||||
layers, attrs
|
||||
)
|
||||
with self.assertRaises(NotImplementedError):
|
||||
runner._get_attn_config()
|
||||
|
||||
def test_sync_writes_full_layers_only(self):
|
||||
runner = _stub_runner(_tiny_gpt_oss_model(), disable_radix_cache=False)
|
||||
# init_cache_pools skips the pool on an SWA model (see
|
||||
# test_sliding_window_model_gets_no_pool), so attach one by hand: the
|
||||
# layer-type filtering in _sync_new_kv_to_pool is what the shared
|
||||
# window-aware SWA pool will rely on, and it must stay correct.
|
||||
self.assertIsNone(runner._attention_kv_pool)
|
||||
runner._attention_kv_pool = MlxAttentionKVPool(
|
||||
pool_size=runner._pool_size + 1,
|
||||
num_layers=runner._cache_layout.num_full_attention_layers,
|
||||
n_kv_heads=2,
|
||||
head_dim=16,
|
||||
dtype=mx.float32,
|
||||
)
|
||||
cache = runner._new_native_cache()
|
||||
per_layer_k = {}
|
||||
for layer_idx in range(4):
|
||||
k = mx.full((1, 2, 5, 16), float(layer_idx + 1))
|
||||
cache[layer_idx].update_and_fetch(k, -k)
|
||||
per_layer_k[layer_idx] = k
|
||||
slot_ids = [7, 9, 11]
|
||||
runner._sync_new_kv_to_pool(cache, cache_start=2, slot_ids=slot_ids)
|
||||
for layer_idx, pool_idx in ((1, 0), (3, 1)):
|
||||
got_k, got_v = runner._attention_kv_pool.get_kv(
|
||||
pool_idx, mx.array(slot_ids, dtype=mx.int32)
|
||||
)
|
||||
want = per_layer_k[layer_idx][0, :, 2:5, :].transpose(1, 0, 2)
|
||||
self.assertTrue(mx.array_equal(got_k, want).item())
|
||||
self.assertTrue(mx.array_equal(got_v, -want).item())
|
||||
# Untouched pool slots stay zero (nothing wrote outside the slots).
|
||||
rest_k, _ = runner._attention_kv_pool.get_kv(
|
||||
0, mx.array([1, 2, 3], dtype=mx.int32)
|
||||
)
|
||||
self.assertEqual(mx.abs(rest_k).max().item(), 0.0)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestSwaRadixPath(CustomTestCase):
|
||||
"""Radix-path prefill/decode is token-identical to the per-request path.
|
||||
|
||||
Both runners share one tiny gpt-oss (same random weights). The
|
||||
reference runner uses the ``disable_radix_cache`` per-request path
|
||||
pinned by test_windowed_kv_cache.py; the radix runner replays the
|
||||
same requests through pool sync, prefix hits, and prefix recomputes.
|
||||
"""
|
||||
|
||||
DECODE_STEPS = 6
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
mx.random.seed(7)
|
||||
cls.model = _tiny_gpt_oss_model()
|
||||
patch_model_attention(cls.model)
|
||||
|
||||
def setUp(self):
|
||||
self.reference = _stub_runner(self.model, disable_radix_cache=True)
|
||||
self.radix = _stub_runner(self.model, disable_radix_cache=False)
|
||||
|
||||
def _greedy(self, runner, rid, full_ids, new_ids, prefix_slots, new_slots):
|
||||
tokens = [
|
||||
runner.prefill(
|
||||
req_id=rid,
|
||||
new_token_ids=list(new_ids),
|
||||
full_token_ids=list(full_ids),
|
||||
prefix_slot_ids=list(prefix_slots),
|
||||
new_slot_ids=list(new_slots),
|
||||
req_pool_idx=0,
|
||||
)
|
||||
]
|
||||
for _ in range(self.DECODE_STEPS):
|
||||
tokens.extend(runner.decode_batch([rid]))
|
||||
return tokens
|
||||
|
||||
def _reference_stream(self, prompt):
|
||||
tokens = self._greedy(
|
||||
self.reference, "ref", prompt, prompt, prefix_slots=(), new_slots=()
|
||||
)
|
||||
self.reference.remove_request("ref")
|
||||
return tokens
|
||||
|
||||
def _seed(self, prompt, slots):
|
||||
self._greedy(self.radix, "seed", prompt, prompt, (), slots)
|
||||
self.radix.remove_request("seed")
|
||||
|
||||
def _assert_windowed_bounded(self, rid):
|
||||
cache = self.radix._req_caches[rid]
|
||||
for layer_idx in self.radix._cache_layout.swa_attention_layer_indices:
|
||||
entry = cache[layer_idx]
|
||||
self.assertIsInstance(entry, WindowedAttentionKVCache)
|
||||
self.assertLessEqual(entry.get_kv()[0].shape[2], 2 * TINY_WINDOW)
|
||||
|
||||
def test_cold_prefill_matches_reference(self):
|
||||
prompt = [(i * 7 + 3) % 128 for i in range(20)]
|
||||
want = self._reference_stream(prompt)
|
||||
got = self._greedy(
|
||||
self.radix, "cold", prompt, prompt, (), range(1, len(prompt) + 1)
|
||||
)
|
||||
self.assertEqual(got, want)
|
||||
self._assert_windowed_bounded("cold")
|
||||
|
||||
def test_partial_prefix_hit_recomputes_exactly(self):
|
||||
# Prefix (20) is well past the window (8): the hit must recompute the
|
||||
# whole prefix rather than gather it, a chunked extend continues on
|
||||
# top, and the stream must match one cold reference over the same
|
||||
# tokens with every cache left at the unclamped absolute position.
|
||||
prefix = [(i * 7 + 3) % 128 for i in range(20)]
|
||||
chunk_a, chunk_b = [9, 42, 77, 5], [11, 13, 17]
|
||||
prefix_slots = list(range(1, len(prefix) + 1))
|
||||
self._seed(prefix, prefix_slots)
|
||||
want = self._reference_stream(prefix + chunk_a + chunk_b)
|
||||
|
||||
gathers = []
|
||||
self.radix._cache_with_pool_backed_attention = lambda slots, n: gathers.append(
|
||||
n
|
||||
)
|
||||
self.radix.prefill(
|
||||
req_id="hit",
|
||||
new_token_ids=chunk_a,
|
||||
full_token_ids=prefix + chunk_a,
|
||||
prefix_slot_ids=prefix_slots,
|
||||
new_slot_ids=list(range(30, 34)),
|
||||
req_pool_idx=0,
|
||||
)
|
||||
self.assertEqual(gathers, [], "SWA prefix hits must recompute, not gather")
|
||||
|
||||
got = [self.radix.extend("hit", chunk_b, list(range(34, 37)))]
|
||||
for _ in range(self.DECODE_STEPS):
|
||||
got.extend(self.radix.decode_batch(["hit"]))
|
||||
self.assertEqual(got, want)
|
||||
self._assert_windowed_bounded("hit")
|
||||
expected = len(prefix + chunk_a + chunk_b) + self.DECODE_STEPS
|
||||
for layer_idx in range(4):
|
||||
self.assertEqual(self.radix._req_caches["hit"][layer_idx].offset, expected)
|
||||
|
||||
def test_full_prefix_hit_without_new_tokens(self):
|
||||
# An exact hit leaves no extend tokens: the prefix rebuild supplies
|
||||
# run tokens ending on the last prefix token, whose logits predict
|
||||
# the next token.
|
||||
prompt = [(i * 5 + 11) % 128 for i in range(20)]
|
||||
prefix_slots = list(range(1, len(prompt) + 1))
|
||||
self._seed(prompt, prefix_slots)
|
||||
want = self._reference_stream(prompt)
|
||||
got = self._greedy(self.radix, "exact", prompt, [], prefix_slots, ())
|
||||
self.assertEqual(got, want)
|
||||
|
||||
def test_fused_aot_kernel_serves_full_layers_by_full_pool_index(self):
|
||||
# The fused RoPE+pool-scatter kernel must skip sliding-window layers
|
||||
# and address the pool by the full-attention index, not the cache one.
|
||||
sliding_wrapper = self.model.model.layers[0].self_attn
|
||||
full_wrapper = self.model.model.layers[1].self_attn
|
||||
|
||||
recorded = []
|
||||
original = MLXAttentionWrapper._rope_custom_aot
|
||||
|
||||
def _recording_rope(queries, keys, values, positions, pool_idx, rope_ctx):
|
||||
recorded.append(pool_idx)
|
||||
return queries, keys
|
||||
|
||||
MLXAttentionWrapper._rope_custom_aot = staticmethod(_recording_rope)
|
||||
try:
|
||||
win = WindowedAttentionKVCache(TINY_WINDOW)
|
||||
contig = ContiguousAttentionKVCache(
|
||||
n_kv_heads=2, head_dim=16, max_seq_len=32, dtype=mx.float32
|
||||
)
|
||||
ctx = BatchedDecodeContext(
|
||||
batch_size=1,
|
||||
seq_lens=[0],
|
||||
attention_layer_caches=[[win], [contig]],
|
||||
attention_pool_index_by_layer={0: 0, 1: 1},
|
||||
full_kv_pool_index_by_layer={1: 0},
|
||||
aot=MlxAOTKernelContext(
|
||||
rope=MlxAOTRoPEContext(kernel=MlxAOTRoPEKernel(), kv_pool=None)
|
||||
),
|
||||
)
|
||||
x = mx.random.normal((1, 1, 64))
|
||||
mx.eval(sliding_wrapper._batched_decode(x, ctx))
|
||||
self.assertEqual(recorded, [], "SWA layer must not hit the fused kernel")
|
||||
mx.eval(full_wrapper._batched_decode(x, ctx))
|
||||
# Cache index for layer 1 is 1; its full-pool index is 0.
|
||||
self.assertEqual(recorded, [0], "full layer needs the full-pool index")
|
||||
finally:
|
||||
MLXAttentionWrapper._rope_custom_aot = original
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
@@ -15,11 +15,15 @@ correct discriminator is ``batch.decoding_reqs``, not the chunk length.
|
||||
|
||||
The routing decision was duplicated across the sync and async paths (the bug
|
||||
therefore existed in both). It now lives in the shared
|
||||
``MlxTpModelWorker._route_extend_request`` helper. These tests cover:
|
||||
``MlxTpModelWorker._route_extend_request`` helper, and the sync entry point
|
||||
launches through the async one rather than re-implementing it. These tests
|
||||
cover:
|
||||
|
||||
* the helper decision directly (both paths delegate to it);
|
||||
* the sync wiring, by driving ``_forward_batch_generation_mlx``;
|
||||
* the async wiring, by driving ``_async_extend_batch``.
|
||||
* the helper decision directly;
|
||||
* the async wiring, by driving ``_async_extend_batch``;
|
||||
* the sync entry point, by driving ``_forward_batch_generation_mlx`` --
|
||||
which also guards the delegation, since a divergence there would show up
|
||||
as a routing or token-ordering difference between the two.
|
||||
|
||||
They mock the MLX runner and load no model. Apple-Silicon-only because
|
||||
``tp_worker`` imports ``mlx.core`` at module load.
|
||||
@@ -35,7 +39,9 @@ from types import SimpleNamespace
|
||||
import torch
|
||||
|
||||
from sglang.srt.model_executor.forward_batch_info import ForwardMode
|
||||
from sglang.srt.runtime_context import get_context
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
# CPU marker is AST-parsed "this test exists"; actual CPU-side execution is
|
||||
# gated by the @skipUnless guard below. MLX marker runs for real on the MLX
|
||||
@@ -49,11 +55,15 @@ _SKIP_REASON = "Apple-Silicon-only (tp_worker imports mlx.core at module load)"
|
||||
|
||||
|
||||
class _FakeRunner:
|
||||
"""Records which routing path each request took (sync + async surfaces)."""
|
||||
"""Records which routing path each request took (both worker paths
|
||||
drive the runner through the same start/finalize surface)."""
|
||||
|
||||
def __init__(self, known_rids):
|
||||
self._known = set(known_rids)
|
||||
self.calls: list[tuple[str, str]] = [] # (op, rid)
|
||||
# (op, rid) -> needs_logits as received; guards the worker's
|
||||
# chunk-finality derivation reaching the runner intact.
|
||||
self.logits_flags: dict[tuple[str, str], bool] = {}
|
||||
self._req_caches: dict[str, list] = {}
|
||||
self._counter = 0
|
||||
|
||||
@@ -73,37 +83,27 @@ class _FakeRunner:
|
||||
|
||||
return SimpleNamespace(state=[mx.array([0.0], dtype=mx.float32)])
|
||||
|
||||
# --- sync surface ---
|
||||
def extend(self, rid, new_token_ids, new_slot_ids):
|
||||
self.calls.append(("extend", rid))
|
||||
self._counter += 1
|
||||
return 1000 + self._counter
|
||||
|
||||
def decode_batch(self, rids):
|
||||
for rid in rids:
|
||||
self.calls.append(("decode", rid))
|
||||
return [2000 + i for i in range(len(rids))]
|
||||
|
||||
def prefill(
|
||||
# --- start/finalize surface (shared by the sync and async worker paths) ---
|
||||
def extend_start(
|
||||
self,
|
||||
req_id,
|
||||
new_token_ids,
|
||||
full_token_ids,
|
||||
prefix_slot_ids,
|
||||
new_slot_ids,
|
||||
req_pool_idx,
|
||||
req=None,
|
||||
needs_logits=True,
|
||||
logit_edit_row=None,
|
||||
logprob_spec=None,
|
||||
):
|
||||
self.calls.append(("prefill", req_id))
|
||||
return 3000
|
||||
|
||||
# --- async surface ---
|
||||
def extend_start(self, req_id, new_token_ids, new_slot_ids):
|
||||
import mlx.core as mx
|
||||
|
||||
self.calls.append(("extend_start", req_id))
|
||||
self.logits_flags[("extend_start", req_id)] = needs_logits
|
||||
self._req_caches[req_id] = [self._fake_cache_layer()]
|
||||
return SimpleNamespace(lazy_token=mx.array([0], dtype=mx.int32), req_id=req_id)
|
||||
return SimpleNamespace(
|
||||
lazy_token=mx.array([0], dtype=mx.int32),
|
||||
cache=self._req_caches[req_id],
|
||||
req_id=req_id,
|
||||
lazy_logprobs=None,
|
||||
)
|
||||
|
||||
def prefill_start(
|
||||
self,
|
||||
@@ -114,17 +114,24 @@ class _FakeRunner:
|
||||
new_slot_ids,
|
||||
req_pool_idx,
|
||||
req=None,
|
||||
needs_logits=True,
|
||||
logit_edit_row=None,
|
||||
logprob_spec=None,
|
||||
):
|
||||
import mlx.core as mx
|
||||
|
||||
self.calls.append(("prefill_start", req_id))
|
||||
self.logits_flags[("prefill_start", req_id)] = needs_logits
|
||||
return SimpleNamespace(
|
||||
lazy_token=mx.array([0], dtype=mx.int32),
|
||||
cache=[self._fake_cache_layer()],
|
||||
req_id=req_id,
|
||||
lazy_logprobs=None,
|
||||
)
|
||||
|
||||
def decode_batch_start(self, rids):
|
||||
def decode_batch_start(
|
||||
self, rids, edit_rows=None, logprob_spec=None, logits_hook=None
|
||||
):
|
||||
import mlx.core as mx
|
||||
|
||||
for rid in rids:
|
||||
@@ -133,8 +140,29 @@ class _FakeRunner:
|
||||
lazy_tokens=mx.array([0] * len(rids), dtype=mx.int32),
|
||||
caches=[[self._fake_cache_layer()] for _ in rids],
|
||||
req_ids=list(rids),
|
||||
lazy_logprobs=None,
|
||||
)
|
||||
|
||||
def prefill_finalize(self, pending):
|
||||
return 3000
|
||||
|
||||
def extend_finalize(self, pending):
|
||||
self._counter += 1
|
||||
return 1000 + self._counter
|
||||
|
||||
def decode_batch_finalize(self, pending):
|
||||
return [2000 + i for i in range(len(pending.req_ids))]
|
||||
|
||||
def collect_logprobs(self, lazy_logprobs):
|
||||
return None
|
||||
|
||||
def eval_pending(self, pending):
|
||||
pass
|
||||
|
||||
@staticmethod
|
||||
def cache_state_arrays(caches):
|
||||
return [s for cache_list in caches for c in cache_list for s in c.state]
|
||||
|
||||
|
||||
class _FakeReq:
|
||||
def __init__(self, rid, req_pool_idx=0):
|
||||
@@ -142,6 +170,11 @@ class _FakeReq:
|
||||
self.prefix_indices = torch.empty(0, dtype=torch.long)
|
||||
self.fill_ids = [0]
|
||||
self.req_pool_idx = req_pool_idx
|
||||
# Mirrors Req's chunk-finality contract read by
|
||||
# MlxTpModelWorker._chunk_needs_logits: extend_range=None means
|
||||
# "not truncated" (final chunk / plain prefill).
|
||||
self.extend_range = None
|
||||
self.full_untruncated_fill_ids = self.fill_ids
|
||||
|
||||
def get_fill_ids(self):
|
||||
return self.fill_ids
|
||||
@@ -154,15 +187,26 @@ class _FakeBatch:
|
||||
self.reqs = reqs
|
||||
self.extend_lens = list(extend_lens)
|
||||
self.decoding_reqs = decoding_reqs
|
||||
self.sampling_info = None
|
||||
self.return_logprob = False
|
||||
# Arbitrary but correctly-sized token / slot arrays.
|
||||
self.input_ids = torch.arange(total, dtype=torch.long)
|
||||
self.out_cache_loc = torch.arange(total, dtype=torch.long)
|
||||
|
||||
|
||||
@unittest.skipUnless(_IS_APPLE_SILICON and _HAS_MLX, _SKIP_REASON)
|
||||
class TestMlxExtendRouting(unittest.TestCase):
|
||||
class TestMlxExtendRouting(CustomTestCase):
|
||||
"""Routing contract for MlxTpModelWorker: shared helper + sync + async."""
|
||||
|
||||
@classmethod
|
||||
def setUpClass(cls):
|
||||
# The worker reads --mlx-enable-sampling off the device config bag,
|
||||
# which fails closed before a publish. Routing itself is orthogonal
|
||||
# to sampling, so pin it off for the whole case.
|
||||
cls._config = get_context().override_server_args(mlx_enable_sampling=False)
|
||||
cls._config.install()
|
||||
cls.addClassCleanup(cls._config.restore)
|
||||
|
||||
@staticmethod
|
||||
def _worker(known_rids):
|
||||
from sglang.srt.hardware_backend.mlx.tp_worker import MlxTpModelWorker
|
||||
@@ -170,6 +214,10 @@ class TestMlxExtendRouting(unittest.TestCase):
|
||||
worker = MlxTpModelWorker.__new__(MlxTpModelWorker)
|
||||
worker._mlx_runner = _FakeRunner(known_rids)
|
||||
worker._mlx_active_rids = set()
|
||||
# The sync entry point delegates to the async launch, which guards
|
||||
# pool creation behind this flag; forward_batch_generation has
|
||||
# already run it for real by the time either path is reached.
|
||||
worker._mlx_pool_initialized = True
|
||||
return worker
|
||||
|
||||
# ---------- the shared decision helper ----------
|
||||
@@ -200,13 +248,28 @@ class TestMlxExtendRouting(unittest.TestCase):
|
||||
def test_sync_one_token_continuation_routes_to_extend(self):
|
||||
"""THE REGRESSION (sync): a 1-token continuation must extend, not decode."""
|
||||
runner = self._run_sync([_FakeReq("r1")], [1], {"r1"}, None, ForwardMode.EXTEND)
|
||||
self.assertEqual(runner.ops_for("r1"), ["extend"])
|
||||
self.assertEqual(runner.ops_for("r1"), ["extend_start"])
|
||||
# Untruncated (extend_range None) => final chunk => logits required.
|
||||
self.assertIs(runner.logits_flags[("extend_start", "r1")], True)
|
||||
|
||||
def test_sync_non_final_chunk_skips_logits(self):
|
||||
"""Head-skip derivation: a scheduler-truncated chunk (extend_range.end
|
||||
below the request's full untruncated length) reaches the runner with
|
||||
needs_logits=False; its next-token output is popped as the stale
|
||||
intermediate token, so computing the vocab head for it is pure waste.
|
||||
Everything else about routing is unchanged."""
|
||||
req = _FakeReq("r1")
|
||||
req.full_untruncated_fill_ids = list(range(8))
|
||||
req.extend_range = SimpleNamespace(start=0, end=4) # 4 < 8: non-final
|
||||
runner = self._run_sync([req], [4], {"r1"}, None, ForwardMode.EXTEND)
|
||||
self.assertEqual(runner.ops_for("r1"), ["extend_start"])
|
||||
self.assertIs(runner.logits_flags[("extend_start", "r1")], False)
|
||||
|
||||
def test_sync_genuine_mixed_decode_routes_to_decode(self):
|
||||
p, d = _FakeReq("p1"), _FakeReq("d1")
|
||||
runner = self._run_sync([p, d], [4, 1], {"d1"}, [d], ForwardMode.MIXED)
|
||||
self.assertEqual(runner.ops_for("p1"), ["prefill"])
|
||||
self.assertEqual(runner.ops_for("d1"), ["decode"])
|
||||
self.assertEqual(runner.ops_for("p1"), ["prefill_start"])
|
||||
self.assertEqual(runner.ops_for("d1"), ["decode_start"])
|
||||
|
||||
# ---------- async path: _async_extend_batch ----------
|
||||
|
||||
@@ -216,26 +279,34 @@ class TestMlxExtendRouting(unittest.TestCase):
|
||||
worker = MlxTpModelWorker.__new__(MlxTpModelWorker)
|
||||
worker._mlx_runner = _FakeRunner(known_rids)
|
||||
batch = _FakeBatch(forward_mode, reqs, extend_lens, decoding_reqs)
|
||||
# returns (lazy_stacked, pending_prefills, pending_extends,
|
||||
# pending_mixed_decode, mode)
|
||||
result = worker._async_extend_batch(batch)
|
||||
return worker._mlx_runner, result
|
||||
launch = worker._async_extend_batch(batch)
|
||||
return worker._mlx_runner, launch
|
||||
|
||||
def test_async_one_token_continuation_routes_to_extend(self):
|
||||
"""THE REGRESSION (async): a 1-token continuation must extend, not decode."""
|
||||
runner, result = self._run_async(
|
||||
runner, launch = self._run_async(
|
||||
[_FakeReq("r1")], [1], {"r1"}, None, ForwardMode.EXTEND
|
||||
)
|
||||
self.assertEqual(runner.ops_for("r1"), ["extend_start"])
|
||||
self.assertEqual(len(result[2]), 1) # one pending extend
|
||||
self.assertIsNone(result[3]) # no mixed decode
|
||||
self.assertIs(runner.logits_flags[("extend_start", "r1")], True)
|
||||
self.assertEqual(len(launch.extends), 1) # one pending extend
|
||||
self.assertIsNone(launch.decode) # no mixed decode
|
||||
|
||||
def test_async_non_final_chunk_skips_logits(self):
|
||||
"""Async twin of the head-skip derivation guard."""
|
||||
req = _FakeReq("r1")
|
||||
req.full_untruncated_fill_ids = list(range(8))
|
||||
req.extend_range = SimpleNamespace(start=0, end=4)
|
||||
runner, _ = self._run_async([req], [4], {"r1"}, None, ForwardMode.EXTEND)
|
||||
self.assertEqual(runner.ops_for("r1"), ["extend_start"])
|
||||
self.assertIs(runner.logits_flags[("extend_start", "r1")], False)
|
||||
|
||||
def test_async_genuine_mixed_decode_routes_to_decode(self):
|
||||
p, d = _FakeReq("p1"), _FakeReq("d1")
|
||||
runner, result = self._run_async([p, d], [4, 1], {"d1"}, [d], ForwardMode.MIXED)
|
||||
runner, launch = self._run_async([p, d], [4, 1], {"d1"}, [d], ForwardMode.MIXED)
|
||||
self.assertEqual(runner.ops_for("p1"), ["prefill_start"])
|
||||
self.assertEqual(runner.ops_for("d1"), ["decode_start"])
|
||||
self.assertIsNotNone(result[3]) # pending mixed decode present
|
||||
self.assertIsNotNone(launch.decode) # pending mixed decode present
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
|
||||
@@ -0,0 +1,321 @@
|
||||
"""Unit tests for the MLX windowed per-request attention KV cache.
|
||||
|
||||
``WindowedAttentionKVCache`` keeps only the trailing ``window`` tokens of a
|
||||
sliding-window layer. Every level is pinned against the full-history path it
|
||||
replaces: the cache arrays against a ``ContiguousAttentionKVCache`` trailing
|
||||
slice, the container forward against full-history caches, and
|
||||
``MLXAttentionWrapper`` batched decode against the same wrapper driven by
|
||||
full-history caches.
|
||||
|
||||
Sliding-window layers use this storage on both KV paths; how it composes
|
||||
with the shared pool and radix prefix hits is pinned in
|
||||
test_swa_radix_pool.py.
|
||||
"""
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import importlib.util
|
||||
import unittest
|
||||
|
||||
from sglang.test.ci.ci_register import register_cpu_ci, register_mlx_ci
|
||||
from sglang.test.test_utils import CustomTestCase
|
||||
|
||||
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
|
||||
register_mlx_ci(est_time=10, suite="stage-a-unit-test-mlx")
|
||||
|
||||
_HAS_MLX = (
|
||||
importlib.util.find_spec("mlx") is not None
|
||||
and importlib.util.find_spec("mlx_lm") is not None
|
||||
)
|
||||
_SKIP_REASON = "requires mlx + mlx_lm"
|
||||
|
||||
if _HAS_MLX:
|
||||
import mlx.core as mx
|
||||
from mlx_lm.models import gpt_oss
|
||||
from mlx_lm.models.base import create_causal_mask
|
||||
|
||||
from sglang.srt.hardware_backend.mlx.kv_cache import (
|
||||
BatchedDecodeContext,
|
||||
ContiguousAttentionKVCache,
|
||||
MLXAttentionWrapper,
|
||||
WindowedAttentionKVCache,
|
||||
find_attention_layers,
|
||||
get_layer_window_sizes,
|
||||
make_attention_mask,
|
||||
)
|
||||
from sglang.srt.hardware_backend.mlx.kv_cache.layout import MlxModelCacheLayout
|
||||
|
||||
WINDOW = 8
|
||||
HIDDEN, N_KV_HEADS, HEAD_DIM = 64, 2, 16
|
||||
|
||||
|
||||
def _dense_mask(mask, n_queries: int, offset: int):
|
||||
"""Densify the cheap ``"causal"`` / ``None`` mask forms.
|
||||
|
||||
``make_attention_mask`` returns those instead of a materialised band
|
||||
whenever the window cannot bind, so their width lives in the key tensor
|
||||
rather than in the mask. Densifying keeps width and content checkable
|
||||
for both forms.
|
||||
"""
|
||||
if mask is None or isinstance(mask, str):
|
||||
return create_causal_mask(n_queries, offset)
|
||||
return mask
|
||||
|
||||
|
||||
def _tiny_gpt_oss_model():
|
||||
"""Random-weight 4-layer gpt_oss, alternating sliding/full layers."""
|
||||
return gpt_oss.Model(
|
||||
gpt_oss.ModelArgs(
|
||||
num_hidden_layers=4,
|
||||
num_local_experts=8,
|
||||
num_experts_per_tok=2,
|
||||
vocab_size=128,
|
||||
hidden_size=HIDDEN,
|
||||
intermediate_size=64,
|
||||
head_dim=HEAD_DIM,
|
||||
num_attention_heads=4,
|
||||
num_key_value_heads=N_KV_HEADS,
|
||||
sliding_window=WINDOW,
|
||||
)
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestWindowedCacheEquivalence(CustomTestCase):
|
||||
"""Storage equivalence with the full-history trailing slice.
|
||||
|
||||
Both caches receive identical K/V and only copy it, so the comparisons
|
||||
are exact (``mx.array_equal``), not tolerance-based.
|
||||
"""
|
||||
|
||||
W, H, D = 4, 2, 8
|
||||
|
||||
def _kv(self, S):
|
||||
return mx.random.normal((1, self.H, S, self.D))
|
||||
|
||||
def test_chunk_patterns_match_full_trailing_slice(self):
|
||||
"""Windowed storage == trailing slice of full history, everywhere.
|
||||
|
||||
Also pins what the forward pass depends on: the mask built before
|
||||
``update_and_fetch`` is exactly as wide as the keys it then returns,
|
||||
the kept prefix still covers a full window, ``offset`` stays
|
||||
absolute, and the decode buffer stays bounded by ``2 * window``.
|
||||
"""
|
||||
mx.random.seed(0)
|
||||
for chunks in [
|
||||
(3,), # stays inside the window
|
||||
(4,), # lands exactly on the window
|
||||
(5,), # first chunk already crosses the window
|
||||
(6, 3), # second chunk forces prefix normalisation
|
||||
(2, 2, 2, 2), # repeated small chunks
|
||||
(1, 1, 1), # degenerate single-token chunks
|
||||
(10, 1, 7), # chunk larger than 2*window, then mixed
|
||||
]:
|
||||
full = ContiguousAttentionKVCache(max_seq_len=128)
|
||||
win = WindowedAttentionKVCache(self.W)
|
||||
for S in chunks:
|
||||
k, v = self._kv(S), self._kv(S)
|
||||
mask = win.make_mask(S, window_size=self.W) # runs before update
|
||||
fk, fv = full.update_and_fetch(k, v)
|
||||
wk, wv = win.update_and_fetch(k, v)
|
||||
at = f"chunk S={S} of {chunks}"
|
||||
self.assertTrue(mx.array_equal(wk, fk[:, :, -wk.shape[2] :, :]), at)
|
||||
self.assertTrue(mx.array_equal(wv, fv[:, :, -wv.shape[2] :, :]), at)
|
||||
# When the window cannot bind, make_mask returns the cheap
|
||||
# "causal"/None form whose width is implicit in the key
|
||||
# tensor; densify so the invariant stays checkable either way.
|
||||
dense = _dense_mask(mask, S, wk.shape[2] - S)
|
||||
self.assertEqual(dense.shape[-1], wk.shape[2], f"mask width, {at}")
|
||||
self.assertGreaterEqual(
|
||||
wk.shape[2] - S, min(win.offset - S, self.W), f"prefix, {at}"
|
||||
)
|
||||
self.assertEqual(win.offset, full.offset, at)
|
||||
|
||||
for step in range(4 * self.W):
|
||||
k, v = self._kv(1), self._kv(1)
|
||||
full.write_token(k, v)
|
||||
win.write_token(k, v)
|
||||
fk, _ = full.get_kv()
|
||||
wk, _ = win.get_kv()
|
||||
t = min(win.offset, self.W)
|
||||
at = f"decode step {step} after {chunks}"
|
||||
self.assertTrue(mx.array_equal(wk[:, :, -t:, :], fk[:, :, -t:, :]), at)
|
||||
self.assertEqual(win.offset, full.offset, at)
|
||||
self.assertLessEqual(win.keys.shape[2], 2 * self.W, at)
|
||||
|
||||
def test_decode_reallocates_amortised_not_per_token(self):
|
||||
"""Compaction must stay amortised O(1) on both write paths."""
|
||||
for write in ("write_token", "update_and_fetch"):
|
||||
win = WindowedAttentionKVCache(self.W)
|
||||
big = self._kv(5 * self.W)
|
||||
win.update_and_fetch(big, big)
|
||||
buf, reallocs = win.keys, 0
|
||||
for _ in range(10 * self.W):
|
||||
getattr(win, write)(self._kv(1), self._kv(1))
|
||||
if win.keys is not buf:
|
||||
buf, reallocs = win.keys, reallocs + 1
|
||||
self.assertLessEqual(reallocs, 12, f"{write} reallocated {reallocs}x")
|
||||
|
||||
def test_full_context_mask_raises_once_history_is_unservable(self):
|
||||
win = WindowedAttentionKVCache(self.W)
|
||||
win.update_and_fetch(self._kv(3), self._kv(3))
|
||||
self.assertEqual(win.make_mask(2), "causal") # nothing dropped yet
|
||||
# One oversized chunk is enough: the next update normalises the
|
||||
# prefix to the window, so full context is already unservable.
|
||||
win.update_and_fetch(self._kv(10), self._kv(10))
|
||||
with self.assertRaises(RuntimeError):
|
||||
win.make_mask(2)
|
||||
|
||||
def test_reset_keeps_buffers_and_replays(self):
|
||||
win = WindowedAttentionKVCache(self.W)
|
||||
win.update_and_fetch(self._kv(6), self._kv(6))
|
||||
win.reset()
|
||||
self.assertEqual(win.offset, 0)
|
||||
self.assertIsNotNone(win.keys) # buffer kept for reuse
|
||||
k = self._kv(2)
|
||||
out, _ = win.update_and_fetch(k, k)
|
||||
self.assertEqual(out.shape[2], 2) # no stale prefix survived
|
||||
self.assertTrue(mx.array_equal(out, k))
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestWindowedModelForward(CustomTestCase):
|
||||
"""Container path: chunked prefill + greedy decode on a tiny gpt-oss."""
|
||||
|
||||
def test_chunked_prefill_and_greedy_decode_match_full_history(self):
|
||||
mx.random.seed(0)
|
||||
model = _tiny_gpt_oss_model()
|
||||
windows = get_layer_window_sizes(model)
|
||||
self.assertEqual([windows[i] for i in range(4)], [WINDOW, None, WINDOW, None])
|
||||
ids = (mx.arange(20) * 7 + 3) % 128 # 2.5x window
|
||||
split = 12 # second chunk starts beyond the window
|
||||
|
||||
tokens = {}
|
||||
for name in ("windowed", "full"):
|
||||
cache = [
|
||||
(
|
||||
WindowedAttentionKVCache(windows[i])
|
||||
if name == "windowed" and windows[i] is not None
|
||||
else ContiguousAttentionKVCache(max_seq_len=64)
|
||||
)
|
||||
for i in range(4)
|
||||
]
|
||||
model(ids[None, :split], cache=cache)
|
||||
out = model(ids[None, split:], cache=cache)
|
||||
seq = []
|
||||
for _ in range(2 * WINDOW): # crosses the compaction boundary
|
||||
token = mx.argmax(out[:, -1, :], axis=-1)
|
||||
seq.append(token.item())
|
||||
out = model(token[None], cache=cache)
|
||||
tokens[name] = seq
|
||||
|
||||
self.assertEqual(
|
||||
tokens["windowed"],
|
||||
tokens["full"],
|
||||
"windowed caches diverge from full history",
|
||||
)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestWindowedBatchedDecode(CustomTestCase):
|
||||
"""The production decode path (``MLXAttentionWrapper._batched_decode``).
|
||||
|
||||
Windowed and full-history caches are driven through the *same* wrapper
|
||||
with the same inputs. The wrapper slices the trailing window off
|
||||
whatever ``get_kv`` returns, so both runs must build byte-identical SDPA
|
||||
inputs and the outputs must be bit-equal, not merely close.
|
||||
"""
|
||||
|
||||
def test_chained_decode_across_compaction_boundary(self):
|
||||
"""Decode steps built in still-lazy pairs, riding a compaction.
|
||||
|
||||
Pairs mirror ``decode_batch_start_chained``: step N+1's graph is
|
||||
built before step N materialises. Compaction allocates a fresh
|
||||
buffer instead of mutating in place, so step N's returned views must
|
||||
stay valid. Prefill lengths 15 and 6 put the oversized-chunk shrink
|
||||
in the first pair and a steady-state rebuild in a later one.
|
||||
"""
|
||||
mx.random.seed(1)
|
||||
attn = _tiny_gpt_oss_model().model.layers[0].self_attn
|
||||
wrapper = MLXAttentionWrapper(attn, layer_idx=0, window_size=WINDOW)
|
||||
wins, fulls = [], []
|
||||
for length in (15, 6):
|
||||
x = mx.random.normal((1, length, HIDDEN))
|
||||
win = WindowedAttentionKVCache(WINDOW)
|
||||
full = ContiguousAttentionKVCache(max_seq_len=64)
|
||||
for cache in (win, full):
|
||||
attn(x, make_attention_mask(length, 0, window_size=WINDOW), cache=cache)
|
||||
wins.append(win)
|
||||
fulls.append(full)
|
||||
|
||||
def build_step(x_step, caches):
|
||||
ctx = BatchedDecodeContext(
|
||||
batch_size=len(caches),
|
||||
seq_lens=[c.offset for c in caches],
|
||||
attention_layer_caches=[caches],
|
||||
)
|
||||
return wrapper._batched_decode(x_step, ctx)
|
||||
|
||||
for pair in range(6):
|
||||
steps = [mx.random.normal((len(wins), 1, HIDDEN)) for _ in range(2)]
|
||||
# Both graphs are built before either materialises.
|
||||
got = [(build_step(x, wins), build_step(x, fulls)) for x in steps]
|
||||
mx.eval(got)
|
||||
for tag, (windowed, full) in zip("ab", got):
|
||||
self.assertTrue(
|
||||
mx.array_equal(windowed, full), f"pair {pair}{tag} diverges"
|
||||
)
|
||||
for win, full in zip(wins, fulls):
|
||||
self.assertEqual(win.offset, full.offset)
|
||||
|
||||
|
||||
@unittest.skipUnless(_HAS_MLX, _SKIP_REASON)
|
||||
class TestModelRunnerCacheWiring(CustomTestCase):
|
||||
"""``_new_native_cache``/``_acquire_cache`` wiring without loading weights."""
|
||||
|
||||
def _stub_runner(self, window_map):
|
||||
from sglang.srt.hardware_backend.mlx.model_runner import MlxModelRunner
|
||||
|
||||
layers, attrs = find_attention_layers(_tiny_gpt_oss_model())
|
||||
runner = MlxModelRunner.__new__(MlxModelRunner)
|
||||
runner._cache_layout = MlxModelCacheLayout.from_attention_discovery(
|
||||
layers, attrs, layer_window_sizes=window_map
|
||||
)
|
||||
runner._max_seq_len = 4096
|
||||
runner._cache_pool = []
|
||||
return runner
|
||||
|
||||
def test_windowed_only_for_sliding_layers_and_reset_on_reuse(self):
|
||||
runner = self._stub_runner(get_layer_window_sizes(_tiny_gpt_oss_model()))
|
||||
cache = runner._new_native_cache()
|
||||
self.assertEqual(
|
||||
[type(c) for c in cache],
|
||||
[
|
||||
WindowedAttentionKVCache,
|
||||
ContiguousAttentionKVCache,
|
||||
WindowedAttentionKVCache,
|
||||
ContiguousAttentionKVCache,
|
||||
],
|
||||
)
|
||||
self.assertEqual(cache[0].window, WINDOW)
|
||||
|
||||
# Models without container windows have an empty map, so every
|
||||
# attention layer keeps a contiguous full-history cache.
|
||||
for c in self._stub_runner({})._new_native_cache():
|
||||
self.assertIsInstance(c, ContiguousAttentionKVCache)
|
||||
|
||||
k = mx.random.normal((1, N_KV_HEADS, 10, HEAD_DIM))
|
||||
cache[0].update_and_fetch(k, k)
|
||||
cache[1].update_and_fetch(k, k)
|
||||
runner._release_cache(cache)
|
||||
reused = runner._acquire_cache()
|
||||
self.assertIs(reused, cache)
|
||||
for c in reused:
|
||||
self.assertEqual(c.offset, 0)
|
||||
# A stale local buffer would prepend the previous request's KV.
|
||||
out, _ = reused[0].update_and_fetch(k, k)
|
||||
self.assertEqual(out.shape[2], 10)
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
unittest.main()
|
||||
Reference in New Issue
Block a user