[Refactor] Pass GLM DSA top-k IDs explicitly

tests/engine/test_glm52_moe_train_engine.py

@@ -204,7 +204,6 @@ def test_sp2_ep4_micro2_compile_offload_train_step(self):
         engine.init_model_weights()
         sp_mesh = init_data_mesh(str(DEVICE), sp_size=2)["sp"]
         data_batches = []
-        seq_ctx_list = []
 
         try:
             for micro_batch_idx in range(4):
@@ -214,7 +213,6 @@ def test_sp2_ep4_micro2_compile_offload_train_step(self):
                 data = {"seq_ctx": full_seq_ctx, "shifted_labels": input_ids[:, 1:]}
                 loss_ctx = engine.model.build_loss_ctx_batch([data], sp_mesh=sp_mesh)[0]
                 seq_ctx = full_seq_ctx.split(sp_mesh)
-                seq_ctx_list.append(seq_ctx)
                 data_batches.append(ModelItem(seq_ctx=seq_ctx, loss_ctx=loss_ctx))
 
             with mock.patch.dict(
@@ -230,9 +228,6 @@ def test_sp2_ep4_micro2_compile_offload_train_step(self):
             self.assertTrue(math.isfinite(step_info["logs_info"]["reduced_mtp_loss"]))
             self.assertTrue(math.isfinite(float(grad_norm)))
             self.assertTrue(engine.optimizer.state)
-            for seq_ctx in seq_ctx_list:
-                self.assertEqual(seq_ctx.dsa_topk_cache.indices, {})
-                self.assertEqual(seq_ctx.dsa_topk_cache.offloaded, {})
         finally:
             del engine
             torch.cuda.empty_cache()

tests/model/test_glm52_moe.py

@@ -8,6 +8,8 @@
 TestGlm52RouterBias
     test_scratch_init_zeroes_main_and_mtp_biases: 从头初始化清零主干与 MTP router bias。
     test_update_bias_handles_main_and_shared_mtp_loads: bias 更新覆盖主干并聚合共享 MTP 深度。
+TestGlm52ExplicitDsaDataflow
+    test_model_forward_backward_without_sequence_context_cache: 模型通过显式 IDs 完成前反向。
 TestGlm52SequenceParallel
     test_mtp_loss_and_gradients_match_full_sequence: SP2 的 MTP loss 与梯度匹配完整序列。
 """
@@ -222,6 +224,28 @@ def test_update_bias_handles_main_and_shared_mtp_loads(self):
         )
 
 
+@unittest.skipUnless(torch.cuda.is_available(), "requires CUDA")
+class TestGlm52ExplicitDsaDataflow:
+    def test_model_forward_backward_with_explicit_dsa_dataflow(self):
+        # 验证 GLM public forward/backward 经显式 DSA IDs 数据流产生有限 loss 和梯度。
+        config = _tiny_glm52_config()
+        config.mtp_config = None
+        model = config.build().to(device="cuda", dtype=torch.bfloat16)
+        model.init_weights()
+
+        input_ids = torch.tensor([[2, 3, 4, 5]], device="cuda")
+        shifted_labels = torch.tensor([[3, 4, 5, 6]], device="cuda")
+        seq_ctx = SequenceContext.from_input_ids((input_ids,), device="cuda")
+        data = {"seq_ctx": seq_ctx, "shifted_labels": shifted_labels}
+        loss_ctx = model.build_loss_ctx_batch([data], sp_mesh=None)[0]
+
+        output = model(seq_ctx=seq_ctx, loss_ctx=loss_ctx)
+        output["loss"].backward()
+
+        assert torch.isfinite(output["loss"])
+        assert any(parameter.grad is not None for parameter in model.parameters())
+
+
 @unittest.skipUnless(torch.cuda.device_count() >= 2, "requires 2 CUDA devices")
 class TestGlm52SequenceParallel(DeterministicDDPTestCase):
     def test_mtp_loss_and_gradients_match_full_sequence(self):

tests/module/attention/test_dsa_mla.py

@@ -23,14 +23,13 @@
 import pytest
 import torch
 import torch.distributed as dist
-import torch.nn as nn
 from torch.distributed.algorithms._checkpoint.checkpoint_wrapper import CheckpointImpl
 
 from xtuner._testing import DeterministicDDPTestCase
 from xtuner.v1.data_proto import SequenceContext
 from xtuner.v1.model.utils import checkpoint_wrapper
 from xtuner.v1.module.attention import DSAMLAConfig
-from xtuner.v1.module.attention.dsa_topk_sharing import register_dsa_topk_decoder_lifecycle_hooks
+from xtuner.v1.module.decoder_layer.dense_decoder_layer import DenseDecoderLayer
 from xtuner.v1.ops.sparse_mla import dsa_topk_indices, sparse_mla
 from xtuner.v1.utils.test_utils import init_data_mesh
 
@@ -103,6 +102,10 @@ def _tiny_dsa_attention(
     indexer_types: list[str] | None = None,
     layer_idx: int = 0,
 ):
+    return _tiny_dsa_config(indexer_types).build(hidden_size=4, layer_idx=layer_idx)
+
+
+def _tiny_dsa_config(indexer_types: list[str] | None = None) -> DSAMLAConfig:
     return DSAMLAConfig(
         num_attention_heads=2,
         head_dim=2,
@@ -116,26 +119,17 @@ def _tiny_dsa_attention(
         index_n_heads=2,
         indexer_types=indexer_types,
         sparse_mla_backend="torch",
-    ).build(hidden_size=4, layer_idx=layer_idx)
-
-
-class _TinyDsaDecoderBlock(nn.Module):
-    def __init__(self, attention: nn.Module) -> None:
-        super().__init__()
-        self.self_attn = attention
-        register_dsa_topk_decoder_lifecycle_hooks(self)
-
-    def forward(
-        self,
-        hidden_states: torch.Tensor,
-        position_embeddings: tuple[torch.Tensor, torch.Tensor],
-        seq_ctx: SequenceContext,
-    ) -> torch.Tensor:
-        return self.self_attn(
-            hidden_states=hidden_states,
-            position_embeddings=position_embeddings,
-            seq_ctx=seq_ctx,
-        )["projected_output"]
+    )
+
+
+def _tiny_dsa_decoder(indexer_types: list[str], layer_idx: int) -> DenseDecoderLayer:
+    return DenseDecoderLayer(
+        hidden_size=4,
+        intermediate_size=8,
+        hidden_act="silu",
+        attention_config=_tiny_dsa_config(indexer_types),
+        layer_idx=layer_idx,
+    )
 
 
 class TestTorchSparseMLA:
@@ -185,55 +179,69 @@ def test_packed_inputs_respect_causal_boundaries_and_backward(self):
         assert outputs["raw_output"].shape == (1, 5, 6)
         assert torch.isfinite(outputs["projected_output"]).all()
         assert torch.isfinite(hidden_states.grad).all()
-        topk = seq_ctx.dsa_topk_cache.indices[0]
+        topk = outputs["dsa_topk_ids"]
+        assert topk.dtype == torch.int32
+        assert topk.is_contiguous()
         for token_idx, seq_start in [(0, 0), (1, 0), (2, 2), (3, 2), (4, 2)]:
             valid_indices = topk[token_idx, 0][topk[token_idx, 0] != -1]
             assert valid_indices.numel() == token_idx - seq_start + 1
             assert valid_indices.min().item() >= seq_start
             assert valid_indices.max().item() <= token_idx
 
-    def test_shared_layers_reuse_topk_without_cross_context_leak(self):
-        # 验证 shared attention 复用同一 SequenceContext 的 source top-k，其他 context 保持独立。
+    def test_shared_layer_consumes_explicit_topk_ids(self):
+        # 验证 shared attention 复用显式 IDs，并在漏传时立即报错。
         torch.manual_seed(0)
         source_attention = _tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=0)
         shared_attention = _tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=1)
         position_embeddings = (torch.ones(1, 4, 2), torch.zeros(1, 4, 2))
         hidden_states = torch.randn(1, 4, 4)
         seq_ctx = SequenceContext.from_input_ids((torch.tensor([[1, 2, 3, 4]]),), device="cpu")
 
-        source_attention(hidden_states, position_embeddings, seq_ctx)
-        source_topk = seq_ctx.dsa_topk_cache.indices[0]
-        shared_output = shared_attention(hidden_states, position_embeddings, seq_ctx)["projected_output"]
-
-        other_seq_ctx = SequenceContext.from_input_ids((torch.tensor([[5, 6, 7, 8]]),), device="cpu")
-        source_attention(torch.randn(1, 4, 4), position_embeddings, other_seq_ctx)
+        source_outputs = source_attention(hidden_states, position_embeddings, seq_ctx)
+        dsa_topk_ids = source_outputs["dsa_topk_ids"]
+        shared_outputs = shared_attention(
+            hidden_states,
+            position_embeddings,
+            seq_ctx,
+            dsa_topk_ids=dsa_topk_ids,
+        )
 
-        assert torch.isfinite(shared_output).all()
-        assert seq_ctx.dsa_topk_cache.indices[0] is source_topk
-        assert other_seq_ctx.dsa_topk_cache.indices[0] is not source_topk
+        assert torch.isfinite(shared_outputs["projected_output"]).all()
+        assert shared_outputs["dsa_topk_ids"] is dsa_topk_ids
+        with pytest.raises(RuntimeError, match="requires dsa_topk_ids"):
+            shared_attention(hidden_states, position_embeddings, seq_ctx)
 
-    def test_reentrant_checkpoint_reuses_and_releases_topk(self):
-        # 验证真实 source/shared decoder 经 reentrant checkpoint 重算后梯度有限且缓存释放。
+    def test_reentrant_checkpoint_preserves_explicit_topk_storage(self):
+        # 验证真实 decoder 的 flat IDs 输入输出可经 reentrant checkpoint 完成反向。
         torch.manual_seed(0)
         source_block = checkpoint_wrapper(
-            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=0)),
+            _tiny_dsa_decoder(["full", "shared"], layer_idx=0),
             checkpoint_impl=CheckpointImpl.REENTRANT,
         )
         shared_block = checkpoint_wrapper(
-            _TinyDsaDecoderBlock(_tiny_dsa_attention(indexer_types=["full", "shared"], layer_idx=1)),
+            _tiny_dsa_decoder(["full", "shared"], layer_idx=1),
             checkpoint_impl=CheckpointImpl.REENTRANT,
         )
         hidden_states = torch.randn(1, 4, 4, requires_grad=True)
         position_embeddings = (torch.ones(1, 4, 2), torch.zeros(1, 4, 2))
         seq_ctx = SequenceContext.from_input_ids((torch.tensor([[1, 2, 3, 4]]),), device="cpu")
 
-        output = source_block(hidden_states, position_embeddings=position_embeddings, seq_ctx=seq_ctx)
-        output = shared_block(output, position_embeddings=position_embeddings, seq_ctx=seq_ctx)
-        output.square().mean().backward()
+        source_hidden, source_ids = source_block(
+            hidden_states,
+            position_embeddings=position_embeddings,
+            seq_ctx=seq_ctx,
+        )
+        shared_hidden, shared_ids = shared_block(
+            source_hidden,
+            source_ids,
+            position_embeddings=position_embeddings,
+            seq_ctx=seq_ctx,
+        )
+        assert shared_ids.untyped_storage().data_ptr() == source_ids.untyped_storage().data_ptr()
+        shared_hidden.square().mean().backward()
 
         assert torch.isfinite(hidden_states.grad).all()
-        assert seq_ctx.dsa_topk_cache.indices == {}
-        assert seq_ctx.dsa_topk_cache.offloaded == {}
+        assert source_ids.dtype == torch.int32
 
 
 class TestAcceleratedSparseMLA:
@@ -322,12 +330,13 @@ def test_packed_attention_matches_full_sequence(self):
         full_output_grad = torch.randn(1, 8, 4, device="cuda")
 
         full_seq_ctx = SequenceContext.from_input_ids(packed_input_ids, device="cuda")
-        expected_output = attention(
+        expected_outputs = attention(
             full_hidden_states,
             position_embeddings=full_position_embeddings,
             seq_ctx=full_seq_ctx,
-        )["projected_output"]
-        expected_topk = full_seq_ctx.dsa_topk_cache.indices[0].clone()
+        )
+        expected_output = expected_outputs["projected_output"]
+        expected_topk = expected_outputs["dsa_topk_ids"].clone()
         expected_output.backward(full_output_grad)
         expected_input_grad = full_hidden_states.grad.clone()
         attention.zero_grad(set_to_none=True)
@@ -338,12 +347,13 @@ def test_packed_attention_matches_full_sequence(self):
         shard_start = sp_seq_ctx.sp_rank * shard_size
         shard_end = shard_start + shard_size
         local_hidden_states = full_hidden_states.detach()[:, shard_start:shard_end].clone().requires_grad_()
-        local_output = attention(
+        local_outputs = attention(
             local_hidden_states,
             position_embeddings=tuple(x[:, shard_start:shard_end] for x in full_position_embeddings),
             seq_ctx=sp_seq_ctx,
-        )["projected_output"]
-        local_topk = sp_seq_ctx.dsa_topk_cache.indices[0]
+        )
+        local_output = local_outputs["projected_output"]
+        local_topk = local_outputs["dsa_topk_ids"]
         local_output.backward(full_output_grad[:, shard_start:shard_end])
 
         gathered_output = [torch.empty_like(local_output) for _ in range(2)]

tests/module/test_dense_decoder_layer.py

@@ -87,11 +87,20 @@ def test_batched_inputs_match_independent_forwards(self):
         )
 
         assert isinstance(outputs, tuple)
-        for output, reference_output in zip(outputs, reference_outputs):
+        n = len(hidden_states)
+        assert len(outputs) == 2 * n
+        output_hidden = outputs[:n]
+        output_ids = outputs[n:]
+        reference_hidden = tuple(result[0] for result in reference_outputs)
+        reference_ids = tuple(result[1] for result in reference_outputs)
+        for output, reference_output in zip(output_hidden, reference_hidden):
             torch.testing.assert_close(output, reference_output)
+        for dsa_topk_ids, reference_dsa_topk_ids in zip(output_ids, reference_ids):
+            torch.testing.assert_close(dsa_topk_ids, reference_dsa_topk_ids)
+            assert dsa_topk_ids.dtype == torch.int32
 
-        sum(output.sum() for output in outputs).backward()
-        sum(output.sum() for output in reference_outputs).backward()
+        sum(output.sum() for output in output_hidden).backward()
+        sum(output.sum() for output in reference_hidden).backward()
 
         for hidden, reference_hidden in zip(hidden_states, reference_hidden_states):
             torch.testing.assert_close(hidden.grad, reference_hidden.grad)

xtuner/v1/data_proto/__init__.py

@@ -1,7 +1,6 @@
-from .sequence_context import DSATopKCacheState, SequenceContext
+from .sequence_context import SequenceContext
 
 
 __all__ = [
-    "DSATopKCacheState",
     "SequenceContext",
 ]

xtuner/v1/data_proto/sequence_context.py

@@ -1,5 +1,4 @@
 # Copyright (c) OpenMMLab. All rights reserved.
-import itertools
 from typing import cast
 
 import torch
@@ -9,52 +8,6 @@
 from .utils import gather_for_sequence_parallel, pad_to_multiple_of, split_for_sequence_parallel
 
 
-_DSA_TOPK_CONTEXT_IDS = itertools.count()
-
-
-class DSATopKCacheState:
-    """Mutable DSA cross-layer top-k cache, scoped to one microbatch.
-
-    For example, if source layer 2 provides top-k indices to layers 2, 3, and 4,
-    its original forward stores ``indices[2]``. After layer 4's no-grad
-    checkpoint forward, ``checkpoint_active`` becomes true and top-k offload may
-    replace that entry with ``offloaded[2]``. Backward then replays layers 4, 3,
-    and 2; layer 2 removes the cache and adds 2 to ``released_sources``. If one
-    physical MTP source is reused at two logical depths, both MTP counters start
-    at 2 so the cache is transferred and released only after the second use in
-    each phase.
-    """
-
-    indices: dict[int, torch.Tensor]  # GPU-resident top-k, keyed by source layer.
-    offloaded: dict[int, str]  # OffloadManager key for each CPU-resident source.
-    released_sources: set[int]  # Sources whose backward replay lifetime has ended.
-    checkpoint_active: bool  # Whether checkpoint forward retained this cache for replay.
-    context_id: int  # Process-local identifier used to make offload keys unique.
-    mtp_forward_uses_remaining: dict[int, int]  # Original-forward MTP uses left per shared source.
-    mtp_replays_remaining: dict[int, int]  # Backward MTP replays left per shared source.
-
-    def __init__(
-        self,
-        *,
-        indices: dict[int, torch.Tensor] | None = None,
-        offloaded: dict[int, str] | None = None,
-        released_sources: set[int] | None = None,
-        checkpoint_active: bool = False,
-        context_id: int | None = None,
-        mtp_forward_uses_remaining: dict[int, int] | None = None,
-        mtp_replays_remaining: dict[int, int] | None = None,
-    ) -> None:
-        # topk_indices format: {source_layer_idx: [seq_len, kv_group, topk]}.
-        # Invalid/padded sparse slots are represented by -1.
-        self.indices = {} if indices is None else indices
-        self.offloaded = {} if offloaded is None else offloaded
-        self.released_sources = set() if released_sources is None else released_sources
-        self.checkpoint_active = checkpoint_active
-        self.context_id = next(_DSA_TOPK_CONTEXT_IDS) if context_id is None else context_id
-        self.mtp_forward_uses_remaining = {} if mtp_forward_uses_remaining is None else mtp_forward_uses_remaining
-        self.mtp_replays_remaining = {} if mtp_replays_remaining is None else mtp_replays_remaining
-
-
 # Avoid using dataclass decorator here to get rid of extra ops called in pytorch 2.8 and above
 # The extra ops is introduced by function _apply_to_tensors in
 # https://github.com/pytorch/pytorch/blob/v2.8.0/torch/distributed/fsdp/_fully_shard/_fsdp_state.py
@@ -97,7 +50,6 @@ class SequenceContext:
     # moe routed_experts
     rollout_routed_experts: torch.Tensor | None
     offload_rollout_routed_experts: bool
-    dsa_topk_cache: DSATopKCacheState
 
     # Private backing attributes for SP shard reconstruction
     _raw_input_ids: torch.LongTensor | None
@@ -127,7 +79,6 @@ def __init__(
         num_img_tokens: list[list[int]] | None = None,
         rollout_routed_experts: torch.Tensor | None = None,
         offload_rollout_routed_experts: bool = False,
-        dsa_topk_cache: DSATopKCacheState | None = None,
         # SP shard metadata: private, accessed via properties below
         raw_input_ids: torch.LongTensor | None = None,
         raw_inputs_embeds: torch.FloatTensor | None = None,
@@ -162,7 +113,6 @@ def __init__(
         self.num_img_tokens = num_img_tokens
         self.rollout_routed_experts = rollout_routed_experts
         self.offload_rollout_routed_experts = offload_rollout_routed_experts
-        self.dsa_topk_cache = DSATopKCacheState() if dsa_topk_cache is None else dsa_topk_cache
         self._raw_input_ids = raw_input_ids
         self._raw_inputs_embeds = raw_inputs_embeds
         self._shard_start = shard_start
@@ -551,7 +501,6 @@ def copy(self, **overrides) -> Self:
             offload_rollout_routed_experts=overrides.get(
                 "offload_rollout_routed_experts", self.offload_rollout_routed_experts
             ),
-            dsa_topk_cache=overrides.get("dsa_topk_cache", self.dsa_topk_cache),
             raw_input_ids=overrides.get("raw_input_ids", self._raw_input_ids),
             raw_inputs_embeds=overrides.get("raw_inputs_embeds", self._raw_inputs_embeds),
             shard_start=overrides.get("shard_start", self._shard_start),
@@ -643,5 +592,4 @@ def data(self) -> dict:
             "num_img_tokens": self.num_img_tokens,
             "rollout_routed_experts": self.rollout_routed_experts,
             "offload_rollout_routed_experts": self.offload_rollout_routed_experts,
-            "dsa_topk_cache": self.dsa_topk_cache,
         }