diff --git a/python/sglang/srt/layers/attention/fla/fused_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_recurrent.py index 67c5859..be3c19c 100644 --- a/python/sglang/srt/layers/attention/fla/fused_recurrent.py +++ b/python/sglang/srt/layers/attention/fla/fused_recurrent.py @@ -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) diff --git a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py index b45738c..375dd6d 100644 --- a/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py +++ b/python/sglang/srt/layers/attention/fla/fused_sigmoid_gating_recurrent.py @@ -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) diff --git a/python/sglang/srt/layers/attention/flashattention_backend.py b/python/sglang/srt/layers/attention/flashattention_backend.py index 3a385db..7eb4233 100644 --- a/python/sglang/srt/layers/attention/flashattention_backend.py +++ b/python/sglang/srt/layers/attention/flashattention_backend.py @@ -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 diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index 8e32640..4c2a84e 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -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. --- /dev/null +++ b/test/registered/unit/mem_cache/test_mamba_prefill_boundary.py @@ -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() --- /dev/null +++ b/test/registered/unit/mem_cache/test_mamba_shared_prefix_lifetime.py @@ -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()