Reduce shared KV materialize synchronization
The shared-KV materialize path was spending time in Python-observed CUDA tensor predicates and dynamic-shape remap helpers. Keep the runtime changes that move the hot paged path to slot-based device remapping, while dropping the NVTX experiment from this commit so profiling annotations do not become part of the runtime surface yet.\n\nThe MLA read path now passes the real page table as the page remap source, which keeps paged topk indices tied to the same logical page-table domain used to build the dense materialized KV view.\n\nConstraint: CP shared KV still needs a dense per-call view before deeper Phase4/Phase5 layout changes remove the materialize cost.\nRejected: Keep NVTX ranges in this commit | user requested reverting NVTX instrumentation before commit\nRejected: Restore compact unique-page remap everywhere | it reintroduces CUDA sync-prone dynamic-shape ops on the hot paged materialize path\nConfidence: medium\nScope-risk: moderate\nDirective: Benchmark slot-remap buffer size against compact unique-page remap before treating this as the final performance path; Phase4/5 should reduce materialize instead of relying on this aggregation path.\nTested: git diff --check on changed files; python -m py_compile on changed runtime/backend/test files; grep confirmed NVTX symbols removed\nNot-tested: pytest blocked locally by missing pybase64 dependency; multi-node PD runtime not rerun
This commit is contained in:
laoyao0822
@@ -366,6 +366,35 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[8:12]))
|
||||
self.assertTrue(torch.equal(dense_kv[20:24], kv_cache[20:24]))
|
||||
|
||||
def test_materialize_token_kv_fast_path_avoids_python_tensor_predicates(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
|
||||
kv_cache = torch.arange(0, 16, dtype=torch.float32).view(16, 1, 1)
|
||||
# Page 1 is owned by CP rank 0, so this also covers the no-local-page
|
||||
# branch without using torch.any(owned_mask) in Python control flow.
|
||||
logical_locs = torch.tensor([4, 5], dtype=torch.int64)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=False
|
||||
), patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
|
||||
), patch.object(
|
||||
runtime.torch, "any", side_effect=AssertionError("torch.any sync")
|
||||
), patch.object(
|
||||
runtime.torch, "equal", side_effect=AssertionError("torch.equal sync")
|
||||
):
|
||||
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_locs.tolist(), [4, 5])
|
||||
self.assertEqual(float(dense_kv.abs().sum().item()), 0.0)
|
||||
|
||||
def test_materialize_token_kv_keeps_dense_shape_for_shared_remap_source(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
@@ -414,7 +443,105 @@ class TestCpSharedKVRuntimeHelpers(unittest.TestCase):
|
||||
page_buffer=page_buffer,
|
||||
logical_pages=logical_pages,
|
||||
layout=layout,
|
||||
)
|
||||
)
|
||||
|
||||
def test_materialize_paged_buffer_fast_path_avoids_python_tensor_predicates(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
|
||||
page_buffer = torch.arange(0, 4 * 3, dtype=torch.uint8).view(4, 3)
|
||||
# Page 1 is owned by CP rank 0, so this also covers the no-local-page
|
||||
# branch without using torch.any(owned_mask) in Python control flow.
|
||||
logical_pages = torch.tensor([1], dtype=torch.int32)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=False
|
||||
), patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
|
||||
), patch.object(
|
||||
runtime.torch, "any", side_effect=AssertionError("torch.any sync")
|
||||
), patch.object(
|
||||
runtime.torch, "equal", side_effect=AssertionError("torch.equal sync")
|
||||
):
|
||||
dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
|
||||
page_buffer=page_buffer,
|
||||
logical_pages=logical_pages,
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_pages.tolist(), [1])
|
||||
self.assertEqual(int(dense_page_buffer.sum().item()), 0)
|
||||
|
||||
def test_materialize_paged_buffer_fast_path_avoids_dynamic_shape_ops(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
|
||||
page_buffer = torch.arange(0, 6 * 3, dtype=torch.uint8).view(6, 3)
|
||||
logical_pages = torch.tensor([1, 2, 5, 6, 0], dtype=torch.int32)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=False
|
||||
), patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
|
||||
), patch.object(
|
||||
runtime.torch, "unique", side_effect=AssertionError("torch.unique sync")
|
||||
), patch.object(
|
||||
runtime.torch, "nonzero", side_effect=AssertionError("torch.nonzero sync")
|
||||
), patch.object(
|
||||
runtime.torch,
|
||||
"searchsorted",
|
||||
side_effect=AssertionError("torch.searchsorted sync"),
|
||||
):
|
||||
dense_page_buffer, dense_pages = runtime.materialize_shared_paged_buffer(
|
||||
page_buffer=page_buffer,
|
||||
logical_pages=logical_pages,
|
||||
layout=layout,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_pages.tolist(), [1, 2, 3, 4, 0])
|
||||
self.assertEqual(list(dense_page_buffer.shape), [6, 3])
|
||||
self.assertTrue(torch.equal(dense_page_buffer[2], page_buffer[1]))
|
||||
self.assertTrue(torch.equal(dense_page_buffer[4], page_buffer[3]))
|
||||
self.assertEqual(int(dense_page_buffer[1].sum().item()), 0)
|
||||
self.assertEqual(int(dense_page_buffer[3].sum().item()), 0)
|
||||
|
||||
def test_materialize_token_kv_page_slot_source_avoids_dynamic_shape_ops(self):
|
||||
from sglang.srt.layers.attention.nsa import cp_shared_kv_runtime as runtime
|
||||
from sglang.srt.mem_cache.cp_shared_kv_layout import CpSharedKVLayout
|
||||
|
||||
layout = CpSharedKVLayout(page_size=4, cp_size=2, cp_rank=1)
|
||||
kv_cache = torch.arange(0, 24, dtype=torch.float32).view(24, 1, 1)
|
||||
logical_locs = torch.tensor([8, 9, 24, 25, -1], dtype=torch.int64)
|
||||
remap_logical_pages = torch.tensor([1, 2, 5, 6], dtype=torch.int32)
|
||||
|
||||
with patch.object(
|
||||
runtime, "cp_shared_kv_debug_enabled", return_value=False
|
||||
), patch.object(
|
||||
runtime, "_all_reduce_materialized_buffer", lambda x, _: x
|
||||
), patch.object(
|
||||
runtime.torch, "unique", side_effect=AssertionError("torch.unique sync")
|
||||
), patch.object(
|
||||
runtime.torch, "nonzero", side_effect=AssertionError("torch.nonzero sync")
|
||||
), patch.object(
|
||||
runtime.torch,
|
||||
"searchsorted",
|
||||
side_effect=AssertionError("torch.searchsorted sync"),
|
||||
):
|
||||
dense_kv, dense_locs = runtime.materialize_shared_token_kv_buffer(
|
||||
kv_cache=kv_cache,
|
||||
logical_locs=logical_locs,
|
||||
remap_logical_pages=remap_logical_pages,
|
||||
layout=layout,
|
||||
page_size=4,
|
||||
)
|
||||
|
||||
self.assertEqual(dense_locs.tolist(), [8, 9, 16, 17, -1])
|
||||
self.assertEqual(list(dense_kv.shape), [20, 1, 1])
|
||||
self.assertTrue(torch.equal(dense_kv[8:12], kv_cache[4:8]))
|
||||
self.assertTrue(torch.equal(dense_kv[16:20], kv_cache[12:16]))
|
||||
self.assertEqual(float(dense_kv[4:8].abs().sum().item()), 0.0)
|
||||
|
||||
|
||||
class TestCpSharedKVLazyDebugLogging(unittest.TestCase):
|
||||
|
||||
Reference in New Issue
Block a user