| |
| |
| |
| |
| @@ -251,7 +251,7 @@ def fused_recurrent_gated_delta_rule_packed_decode_kernel( |
| x = a_val + dt_bias_val |
| softplus_x = tl.where(x <= SOFTPLUS_THRESHOLD, tl.log(1.0 + tl.exp(x)), x) |
| g_val = -tl.exp(A_log_val) * softplus_x |
| - beta_val = tl.sigmoid(b_val).to(b.dtype.element_ty).to(tl.float32) |
| + beta_val = tl.sigmoid(b_val) |
| |
| b_h *= exp(g_val) |
| b_v -= tl.sum(b_h * b_k[None, :], 1) |
| |
| |
| |
| |
| @@ -3,6 +3,7 @@ from typing import Optional |
| import torch |
| import triton |
| import triton.language as tl |
| +import triton.language.extra.libdevice as libdevice |
| |
| |
| @triton.jit(do_not_specialize=["T"]) |
| @@ -184,7 +185,8 @@ def fused_sigmoid_gating_delta_rule_update_kernel( |
| if IS_KDA: |
| b_h *= tl.exp(b_g[:, None]) |
| else: |
| - b_h *= tl.exp(b_g) |
| + # Keep decay rounded before the delta-rule fused update. |
| + b_h = libdevice.mul_rn(b_h, tl.exp(b_g)) |
| |
| # Delta rule: v -= sum(h * k, dim=0) |
| b_v -= tl.sum(b_h * b_k[:, None], 0) |
| @@ -193,7 +195,11 @@ def fused_sigmoid_gating_delta_rule_update_kernel( |
| b_v *= b_beta |
| |
| # Update hidden state: h += k[:, None] * v[None, :] |
| - b_h += b_k[:, None] * b_v[None, :] |
| + if IS_KDA: |
| + b_h += b_k[:, None] * b_v[None, :] |
| + else: |
| + # Match the single-step state write-back arithmetic in verify. |
| + b_h = libdevice.fma_rn(b_k[:, None], b_v[None, :], b_h) |
| |
| # Compute output: o = sum(h * q, dim=0) |
| b_o = tl.sum(b_h * b_q[:, None], 0) |
| |
| |
| |
| |
| @@ -210,12 +210,12 @@ class FlashAttentionBackend(AttentionBackend): |
| # If num_splits == 0, we use a heuristic to automatically determine the number of splits. |
| # We set nums splits to 1 if deterministic inference is enabled. |
| # See https://thinkingmachines.ai/blog/defeating-nondeterminism-in-llm-inference/ for more details. |
| - # Furthermore, FA4 does not support num_splits=0 with CUDA Graph, so we set num_splits to 1 if CUDA Graph is enabled. |
| + # Keep FA3 graph/eager split arithmetic consistent; FA4 also requires this for graphs. |
| self.num_splits = ( |
| 1 |
| if model_runner.server_args.enable_deterministic_inference |
| or ( |
| - self.fa_impl_ver == 4 |
| + self.fa_impl_ver in (3, 4) |
| and not model_runner.server_args.disable_cuda_graph |
| ) |
| else 0 |
| |
| |
| |
| |
| @@ -1467,10 +1467,20 @@ class Scheduler( |
| if self._war_barrier_enabled: |
| self.schedule_stream.wait_stream(self.forward_stream) |
| |
| + # Mamba prefix insertion can free duplicate KV slots and replace |
| + # req_to_token entries with canonical cached slots. Finish that |
| + # transition before preparing a dependent decode/verify batch. |
| + processed_prefill_cache = self._has_pending_mamba_cache_update() |
| + if processed_prefill_cache: |
| + pop_and_process() |
| + |
| # Get the next batch to run |
| batch = self.get_next_batch_to_run() |
| self.cur_batch = batch |
| - disable_overlap_for_batch = self.is_disable_overlap_for_batch(batch) |
| + disable_overlap_for_batch = ( |
| + not processed_prefill_cache |
| + and self.is_disable_overlap_for_batch(batch) |
| + ) |
| |
| # If we do not need to overlap the current batch with the last batch, |
| # we can process the last batch immediately. |
| @@ -1486,7 +1496,7 @@ class Scheduler( |
| |
| # Process the last batch |
| if self.last_batch: |
| - if not disable_overlap_for_batch: |
| + if not disable_overlap_for_batch and not processed_prefill_cache: |
| pop_and_process() |
| elif batch is None: |
| # When the server is idle, do self-check and re-init some states |
| @@ -1503,6 +1513,24 @@ class Scheduler( |
| if envs.SGLANG_ENABLE_STRICT_MEM_CHECK_DURING_BUSY.get(): |
| self.invariant_checker.self_check_during_busy() |
| |
| + def _has_pending_mamba_cache_update(self) -> bool: |
| + # DP/MLP-sync ranks need a collective decision; leave that path unchanged. |
| + if ( |
| + self.require_mlp_sync |
| + or not self.result_queue |
| + or getattr(self.tree_cache, "disable", True) |
| + or not getattr(self.tree_cache, "enable_mamba_extra_buffer", False) |
| + ): |
| + return False |
| + pending_batch = self.result_queue[0][0] |
| + return pending_batch.forward_mode.is_extend() and any( |
| + req.mamba_last_track_seqlen is not None |
| + and not req.finished() |
| + and not req.is_retracted |
| + and not getattr(req, "skip_radix_cache_insert", False) |
| + for req in pending_batch.reqs |
| + ) |
| + |
| def is_disable_overlap_for_batch(self, batch: ScheduleBatch) -> bool: |
| # For two consecutive prefill batches, we disable overlap to improve the TTFT of the first batch. |
| # This might slightly hurt the throughput, so we use an environment variable to control it. |
| |
| |
| @@ -0,0 +1,184 @@ |
| +"""Regression for canonical KV publication before dependent overlap forwards.""" |
| + |
| +import unittest |
| +from collections import deque |
| +from types import SimpleNamespace as NS |
| +from sglang.srt.managers.scheduler import Scheduler |
| + |
| + |
| +class StopLoop(Exception): |
| + pass |
| + |
| + |
| +class Batch: |
| + def __init__(self, name, req): |
| + self.name = name |
| + self.reqs = [req] |
| + self.forward_mode = NS(is_extend=lambda: name == "prefill") |
| + |
| + def copy(self): |
| + return self |
| + |
| + |
| +class Harness: |
| + event_loop_overlap = Scheduler.event_loop_overlap |
| + _has_pending_mamba_cache_update = Scheduler._has_pending_mamba_cache_update |
| + |
| + def __init__(self, finish=False, fail=False): |
| + self.require_mlp_sync = False |
| + self.tree_cache = NS(disable=False, enable_mamba_extra_buffer=True) |
| + self.req = NS( |
| + mamba_last_track_seqlen=384, |
| + is_retracted=False, |
| + skip_radix_cache_insert=False, |
| + finished=lambda: False, |
| + ) |
| + self.last_batch = None |
| + self._engine_paused = False |
| + self._war_barrier_enabled = False |
| + self.is_generation = True |
| + self.events = [] |
| + self.iteration = 0 |
| + self.sent = False |
| + self.done = False |
| + self.finish = finish |
| + self.fail = fail |
| + self.private_slots = [91, 92, 93] |
| + self.canonical_slots = [1, 2, 3] |
| + self.mapping = list(self.private_slots) |
| + self.shared_values = [10, 20, 30] |
| + self.other_request = list(self.canonical_slots) |
| + self.refcount = 1 |
| + self.freed = [] |
| + |
| + def receive(): |
| + self.iteration += 1 |
| + if self.iteration > 3: |
| + raise StopLoop() |
| + return [] |
| + |
| + self.request_receiver = NS(recv_requests=receive) |
| + |
| + def process_input_requests(self, reqs): |
| + pass |
| + |
| + def get_next_batch_to_run(self): |
| + self.events.append("schedule") |
| + if self.done: |
| + return None |
| + if not self.sent: |
| + self.sent = True |
| + return Batch("prefill", self.req) |
| + return Batch("decode", self.req) |
| + |
| + def is_disable_overlap_for_batch(self, batch): |
| + return False |
| + |
| + def run_batch(self, batch): |
| + self.events.append("run:" + batch.name) |
| + if batch.name == "decode": |
| + assert self.mapping == self.canonical_slots, ( |
| + "dependent decode observed unpublished cache mapping" |
| + ) |
| + assert not set(self.mapping) & set(self.freed), ( |
| + "dependent decode references released slots" |
| + ) |
| + return object() |
| + |
| + def process_batch_result(self, batch, result): |
| + self.events.append("process:" + batch.name) |
| + if batch.name == "prefill": |
| + if self.fail: |
| + raise RuntimeError("prefill failed") |
| + if self.finish: |
| + self.done = True |
| + self.freed += self.private_slots |
| + else: |
| + self.freed += self.private_slots |
| + self.mapping = list(self.canonical_slots) |
| + self.refcount += 1 |
| + self.req.mamba_last_track_seqlen = None |
| + else: |
| + self.done = True |
| + self.refcount -= 1 |
| + |
| + def launch_batch_sample_if_needed(self, result): |
| + pass |
| + |
| + def on_idle(self): |
| + pass |
| + |
| + |
| +class TestPrefillBoundary(unittest.TestCase): |
| + def test_publish_before_scheduling_and_consume_once(self): |
| + h = Harness() |
| + with self.assertRaises(StopLoop): |
| + h.event_loop_overlap() |
| + self.assertLess(h.events.index("process:prefill"), h.events.index("run:decode")) |
| + self.assertEqual(h.events.count("process:prefill"), 1) |
| + self.assertEqual(h.events.count("process:decode"), 1) |
| + self.assertEqual(h.freed, [91, 92, 93]) |
| + self.assertEqual(h.refcount, 1) |
| + self.assertEqual(h.other_request, [1, 2, 3]) |
| + self.assertEqual(h.shared_values, [10, 20, 30]) |
| + |
| + def test_finish_in_prefill_is_processed_before_successor_allocation(self): |
| + h = Harness(finish=True) |
| + with self.assertRaises(StopLoop): |
| + h.event_loop_overlap() |
| + self.assertNotIn("run:decode", h.events) |
| + self.assertEqual(h.freed, [91, 92, 93]) |
| + |
| + def test_failed_prefill_cannot_launch_dependent_decode(self): |
| + h = Harness(fail=True) |
| + with self.assertRaisesRegex(RuntimeError, "prefill failed"): |
| + h.event_loop_overlap() |
| + self.assertNotIn("run:decode", h.events) |
| + self.assertEqual(h.freed, []) |
| + |
| + def test_untracked_prefill_retains_overlap(self): |
| + h = Harness() |
| + h.req.mamba_last_track_seqlen = None |
| + h.mapping = list(h.canonical_slots) |
| + with self.assertRaises(StopLoop): |
| + h.event_loop_overlap() |
| + self.assertLess(h.events.index("run:decode"), h.events.index("process:prefill")) |
| + |
| + def test_nonapplicable_paths_do_not_add_barrier(self): |
| + for change in [ |
| + "dp", |
| + "disabled", |
| + "nonmamba", |
| + "no_track", |
| + "finished", |
| + "retracted", |
| + "skip", |
| + "decode", |
| + "empty", |
| + ]: |
| + with self.subTest(change=change): |
| + h = Harness() |
| + h.result_queue = deque([(Batch("prefill", h.req), object())]) |
| + if change == "dp": |
| + h.require_mlp_sync = True |
| + elif change == "disabled": |
| + h.tree_cache.disable = True |
| + elif change == "nonmamba": |
| + h.tree_cache.enable_mamba_extra_buffer = False |
| + elif change == "no_track": |
| + h.req.mamba_last_track_seqlen = None |
| + elif change == "finished": |
| + h.req.finished = lambda: True |
| + elif change == "retracted": |
| + h.req.is_retracted = True |
| + elif change == "skip": |
| + h.req.skip_radix_cache_insert = True |
| + elif change == "decode": |
| + h.result_queue[0] = (Batch("decode", h.req), object()) |
| + elif change == "empty": |
| + h.result_queue.clear() |
| + self.assertFalse(h._has_pending_mamba_cache_update()) |
| + |
| + |
| +if __name__ == "__main__": |
| + unittest.main() |
| |
| |
| @@ -0,0 +1,64 @@ |
| +import unittest |
| +from array import array |
| +import torch |
| +import test_mamba_unittest as mamba_tests |
| +from sglang.srt.mem_cache.base_prefix_cache import ( |
| + InsertParams, |
| + MatchPrefixParams, |
| + EvictParams, |
| +) |
| +from sglang.srt.mem_cache.radix_cache import RadixKey |
| + |
| + |
| +class SharedPins(unittest.TestCase): |
| + def test_shared_prefix_retained_until_both_consumers_release(self): |
| + tree, alloc, pool, make_req = ( |
| + mamba_tests.TestMamba()._setup_tree_and_allocator() |
| + ) |
| + initial = alloc.available_size() |
| + ma = pool.mamba_allocator.available_size() |
| + a = make_req() |
| + av = alloc.alloc(3) |
| + key = RadixKey(array("q", [1, 2, 3])) |
| + tree.insert( |
| + InsertParams(key=key, value=av, mamba_value=a.mamba_pool_idx.unsqueeze(0)) |
| + ) |
| + na = tree.match_prefix(MatchPrefixParams(key=key)).last_device_node |
| + tree.inc_lock_ref(na) |
| + kcache = alloc.get_kvcache().get_key_buffer(3) |
| + marker = torch.full_like(kcache.index_select(0, av), 0.25) |
| + kcache.index_copy_(0, av, marker) |
| + b = make_req() |
| + bv = alloc.alloc(7) |
| + bkey = RadixKey(array("q", [1, 2, 3, 4, 5, 6, 7])) |
| + tree.insert( |
| + InsertParams( |
| + key=bkey, |
| + value=bv, |
| + mamba_value=b.mamba_pool_idx.unsqueeze(0), |
| + prev_prefix_len=0, |
| + ) |
| + ) |
| + nb = tree.match_prefix(MatchPrefixParams(key=bkey)).last_device_node |
| + tree.inc_lock_ref(nb) |
| + free = torch.cat([alloc.free_pages, alloc.release_pages] + alloc.free_group) |
| + self.assertTrue(bool(torch.isin(bv[:3], free).all())) |
| + self.assertFalse(bool(torch.isin(av, free).any())) |
| + self.assertEqual(na.full_lock_ref, 2) |
| + self.assertTrue(torch.equal(na.value, av)) |
| + self.assertTrue(torch.equal(kcache.index_select(0, av), marker)) |
| + self.assertEqual(tree.evict(EvictParams(num_tokens=128)).num_tokens_evicted, 0) |
| + tree.dec_lock_ref(nb) |
| + tree.evict(EvictParams(num_tokens=128)) |
| + self.assertEqual(na.full_lock_ref, 1) |
| + self.assertTrue(torch.equal(kcache.index_select(0, av), marker)) |
| + self.assertEqual(alloc.available_size(), initial - 3) |
| + tree.dec_lock_ref(na) |
| + tree.evict(EvictParams(num_tokens=128)) |
| + tree.sanity_check() |
| + self.assertEqual(alloc.available_size(), initial) |
| + self.assertEqual(pool.mamba_allocator.available_size(), ma) |
| + |
| + |
| +if __name__ == "__main__": |
| + unittest.main() |
|
|
