diff --git a/python/sglang/srt/models/utils.py b/python/sglang/srt/models/utils.py index 96bec941e9bc..16df5d01b60a 100644 --- a/python/sglang/srt/models/utils.py +++ b/python/sglang/srt/models/utils.py @@ -41,7 +41,7 @@ ) from sglang.srt.model_executor.runner import get_is_capture_mode from sglang.srt.model_loader.weight_utils import default_weight_loader -from sglang.srt.runtime_context import get_exec +from sglang.srt.runtime_context import get_exec, get_parallel from sglang.srt.utils import get_current_device_stream_fast, is_cpu, is_cuda, is_hip from sglang.srt.utils.custom_op import register_custom_op @@ -296,6 +296,9 @@ def enable_fused_set_kv_buffer(forward_batch: ForwardBatch): `.view(-> 4D NHD)` reshape, and let the rotary forward pass `flash_layout=False`. See `create_fused_set_kv_buffer_arg` below. """ + # The fused RoPE writer cannot localize DCP slots or filter token owners. + if get_parallel().dcp_enabled: + return False pool = get_token_to_kv_pool() return ( _is_cuda diff --git a/test/registered/unit/models/test_fused_set_kv_buffer_gate.py b/test/registered/unit/models/test_fused_set_kv_buffer_gate.py index c71fafa2f31f..a56aa1251767 100644 --- a/test/registered/unit/models/test_fused_set_kv_buffer_gate.py +++ b/test/registered/unit/models/test_fused_set_kv_buffer_gate.py @@ -95,6 +95,11 @@ def _hip_gate(pool, dcp_kv_mask=None) -> bool: with ( patch.object(model_utils, "_is_hip", True), patch.object(model_utils, "_is_cuda", False), + patch.object( + model_utils, + "get_parallel", + return_value=SimpleNamespace(dcp_enabled=False), + ), patch.object(model_utils, "get_token_to_kv_pool", return_value=pool), ): return model_utils.enable_fused_set_kv_buffer( @@ -156,5 +161,35 @@ def test_decode_context_parallel_mask_declines(self): self.assertFalse(_hip_gate(_static_window_pool(), dcp_kv_mask=object())) +class TestDcpFusedStoreGate(CustomTestCase): + def test_dcp_without_a_batch_mask_uses_the_backend_store(self): + pool = SimpleNamespace( + dtype=torch.bfloat16, + layer_router=SimpleNamespace(sub_pool=lambda name: None), + ) + for cuda in (True, False): + for dcp_size in (1, 2, 4): + with ( + self.subTest(cuda=cuda, dcp_size=dcp_size), + patch.object(model_utils, "_is_cuda", cuda), + patch.object(model_utils, "_is_hip", not cuda), + patch.object(model_utils, "is_cp_active", return_value=False), + patch.object( + model_utils, "get_token_to_kv_pool", return_value=pool + ), + patch.object( + model_utils, + "get_parallel", + return_value=SimpleNamespace(dcp_enabled=dcp_size > 1), + ), + ): + self.assertEqual( + model_utils.enable_fused_set_kv_buffer( + SimpleNamespace(dcp_kv_mask=None) + ), + dcp_size == 1, + ) + + if __name__ == "__main__": unittest.main()