Facade DSA index-cache: MTP topk-reuse state + index-K storage (#28609)

test/registered/unit/layers/attention/test_index_topk_share.py

@@ -0,0 +1,163 @@
+import sys
+from types import SimpleNamespace
+
+import pytest
+import torch
+
+from sglang.srt.layers.attention.index_topk_share import IndexTopKShareState
+from sglang.test.ci.ci_register import register_cpu_ci
+
+register_cpu_ci(est_time=2, suite="base-a-test-cpu")
+
+
+def _batch(
+    reuse: bool, carried, *, is_extend: bool = False, seed_buf=None, seed_select=None
+) -> SimpleNamespace:
+    return SimpleNamespace(
+        reuse_dsa_topk_indices=reuse,
+        forward_mode=SimpleNamespace(
+            is_extend=lambda include_draft_extend_v2: is_extend
+        ),
+        spec_info=SimpleNamespace(
+            dsa_topk_indices=carried,
+            dsa_seed_topk_capture=seed_buf,
+            dsa_seed_topk_select=seed_select,
+        ),
+    )
+
+
+def test_mtp_carry_is_empty_when_reuse_disabled():
+    batch = _batch(reuse=False, carried="old")
+    state = IndexTopKShareState.from_mtp_carry(batch)
+
+    assert state.topk_indices is None
+    state.update("new")
+    state.publish()
+
+    assert state.topk_indices == "new"
+    assert batch.spec_info.dsa_topk_indices == "old"
+
+
+def test_mtp_carry_reads_and_publishes_when_reuse_enabled():
+    batch = _batch(reuse=True, carried="old")
+    state = IndexTopKShareState.from_mtp_carry(batch)
+
+    assert state.topk_indices == "old"
+    state.update("new")
+    state.publish()
+
+    assert state.topk_indices == "new"
+    assert batch.spec_info.dsa_topk_indices == "new"
+
+
+def test_target_carry_stays_local_without_publish():
+    batch = _batch(reuse=False, carried="batch")
+    state = IndexTopKShareState(batch, "layer")
+
+    assert state.topk_indices == "layer"
+    assert batch.spec_info.dsa_topk_indices == "batch"
+
+    state.update(None)
+
+    assert state.topk_indices is None
+    assert batch.spec_info.dsa_topk_indices == "batch"
+
+
+def test_target_none_does_not_fall_back_to_mtp_carry():
+    batch = _batch(reuse=True, carried="batch")
+    state = IndexTopKShareState(batch, None)
+
+    assert state.topk_indices is None
+
+    state.update("layer")
+
+    assert state.topk_indices == "layer"
+    assert batch.spec_info.dsa_topk_indices == "batch"
+
+
+def test_publish_captures_draft_extend_seed():
+    seed_buf = torch.zeros(2, 3, dtype=torch.int64)
+    batch = _batch(reuse=False, carried=None, is_extend=True, seed_buf=seed_buf)
+    state = IndexTopKShareState.from_mtp_carry(batch)
+
+    assert state.should_publish
+    state.update(torch.arange(12, dtype=torch.int64).view(4, 3))
+    state.publish()
+
+    assert torch.equal(seed_buf, torch.arange(6, dtype=torch.int64).view(2, 3))
+    assert batch.spec_info.dsa_topk_indices is None
+
+
+def test_seed_buffer_is_ignored_outside_extend():
+    seed_buf = torch.zeros(2, 3, dtype=torch.int64)
+    batch = _batch(reuse=False, carried=None, is_extend=False, seed_buf=seed_buf)
+    state = IndexTopKShareState.from_mtp_carry(batch)
+
+    assert not state.should_publish
+    state.update(torch.ones(4, 3, dtype=torch.int64))
+    state.publish()
+
+    assert torch.equal(seed_buf, torch.zeros(2, 3, dtype=torch.int64))
+
+
+def test_mtp_iteration_clears_batch_state():
+    batch = _batch(reuse=False, carried="stale")
+
+    with IndexTopKShareState.mtp_iteration(batch) as state:
+        assert state is not None
+        assert batch.reuse_dsa_topk_indices
+        assert batch.spec_info.dsa_topk_indices is None
+        batch.spec_info.dsa_topk_indices = "draft-topk"
+
+    assert not batch.reuse_dsa_topk_indices
+    assert batch.spec_info.dsa_topk_indices is None
+
+
+def test_mtp_iteration_clears_batch_state_on_exception():
+    batch = _batch(reuse=False, carried=None)
+
+    with pytest.raises(RuntimeError, match="draft step blew up"):
+        with IndexTopKShareState.mtp_iteration(batch):
+            batch.spec_info.dsa_topk_indices = "draft-topk"
+            raise RuntimeError("draft step blew up")
+
+    assert not batch.reuse_dsa_topk_indices
+    assert batch.spec_info.dsa_topk_indices is None
+
+
+def test_disabled_mtp_iteration_is_passthrough():
+    batch = _batch(reuse=False, carried="untouched")
+
+    with IndexTopKShareState.mtp_iteration(batch, enabled=False) as state:
+        assert state is None
+        assert not batch.reuse_dsa_topk_indices
+        assert batch.spec_info.dsa_topk_indices == "untouched"
+
+    assert not batch.reuse_dsa_topk_indices
+    assert batch.spec_info.dsa_topk_indices == "untouched"
+
+
+def test_mtp_iteration_preserves_draft_extend_seed():
+    batch = _batch(reuse=False, carried="extend-seed")
+
+    with IndexTopKShareState.mtp_iteration(batch, keep_carry_seed=True) as state:
+        assert state is not None
+        assert batch.spec_info.dsa_topk_indices == "extend-seed"
+        assert state.topk_indices == "extend-seed"
+
+    assert not batch.reuse_dsa_topk_indices
+    assert batch.spec_info.dsa_topk_indices is None
+
+
+def test_mtp_iteration_clears_missing_draft_extend_seed():
+    batch = _batch(reuse=False, carried=None)
+
+    with IndexTopKShareState.mtp_iteration(batch, keep_carry_seed=True):
+        assert batch.spec_info.dsa_topk_indices is None
+
+    assert not batch.reuse_dsa_topk_indices
+    assert batch.spec_info.dsa_topk_indices is None
+
+
+if __name__ == "__main__":
+    sys.exit(pytest.main([__file__]))

test/registered/unit/mem_cache/test_dsa_layer_split_broadcast.py

@@ -104,8 +104,12 @@ def _run(rank: int, world: int, port: int):
     # Indexer buffer owner-broadcast: owner writes a layer-distinct value, then
     # every rank must read it back for every layer.
     for layer_id in range(LAYER_NUM):
+        store_buf = pool.get_index_k_with_scale_buffer(layer_id)
+        assert (
+            store_buf.data_ptr() == pool.index_k_with_scale_buffer[layer_id].data_ptr()
+        )
         if pool._is_layer_owned(layer_id):
-            pool.index_k_with_scale_buffer[layer_id].fill_(layer_id + 10)
+            store_buf.fill_(layer_id + 10)
     torch.cuda.synchronize()
     torch.distributed.barrier()
     for layer_id in range(LAYER_NUM):