Fix DSA top-k pinned buffer growth

ci/scripts/CI_ENV.sh

@@ -28,6 +28,7 @@ export ROLLOUT_DAPO_DATA_PATH=${CI_SHARE_DATA}/rl_test_judger_dapo_math_data.jso
 export GEO_ROLLOUT_DATA_PATH=${CI_SHARE_DATA}/rl_test_judge_geo_data.jsonl
 export TORCH_ALLOW_TF32_CUBLAS_OVERRIDE=0
 export XTUNER_DETERMINISTIC=true
+export TORCHINDUCTOR_DYNAMIC_SCALE_RBLOCK=0
 export XTUNER_USE_LMDEPLOY=1
 export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
 export PYTHONPYCACHEPREFIX=/tmp

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,9 +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.
 
@@ -29,7 +25,7 @@ class DSATopKCacheState:
     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.
+    offload_slot: int  # Stable offload slot among concurrently active microbatches.
     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.
 
@@ -40,7 +36,7 @@ def __init__(
         offloaded: dict[int, str] | None = None,
         released_sources: set[int] | None = None,
         checkpoint_active: bool = False,
-        context_id: int | None = None,
+        offload_slot: int = 0,
         mtp_forward_uses_remaining: dict[int, int] | None = None,
         mtp_replays_remaining: dict[int, int] | None = None,
     ) -> None:
@@ -50,7 +46,7 @@ def __init__(
         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.offload_slot = offload_slot
         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
 

xtuner/v1/model/moe/moe.py

@@ -475,6 +475,14 @@ def post_micro_batch_forward(self, batch_outputs: Sequence[MoEModelOutputs]) ->
         moe_info = cast(MoEBatchForwardInfo, base_info)
         return moe_info
 
+    @staticmethod
+    def _prepare_seq_ctx_topk_cache(seq_ctx_list: Sequence[SequenceContext]) -> None:
+        # A slot owns both runtime offload state and pinned storage. TrainEngine
+        # finishes backward before the next accumulation group, so only contexts
+        # in one model call can be live concurrently and need distinct slots.
+        for offload_slot, seq_ctx in enumerate(seq_ctx_list):
+            seq_ctx.dsa_topk_cache.offload_slot = offload_slot
+
     def _micro_batch_forward(
         self,
         seq_ctx_list: list[SequenceContext],
@@ -486,6 +494,7 @@ def _micro_batch_forward(
         This method processes multiple micro-batches in parallel, similar to how MoEDecoderLayer handles micro-batching
         at the layer level.
         """
+        self._prepare_seq_ctx_topk_cache(seq_ctx_list)
         if self.config.return_hidden_states:
             raise NotImplementedError
 
@@ -624,7 +633,9 @@ def _micro_batch_forward(
                         input_ids=seq_ctx.input_ids.clone() if seq_ctx.input_ids is not None else None,
                         position_ids=seq_ctx.position_ids.clone(),
                         inputs_embeds=seq_ctx.inputs_embeds.clone() if seq_ctx.inputs_embeds is not None else None,
-                        dsa_topk_cache=DSATopKCacheState(),
+                        dsa_topk_cache=DSATopKCacheState(
+                            offload_slot=seq_ctx.dsa_topk_cache.offload_slot,
+                        ),
                     )
                 )
 
@@ -747,6 +758,7 @@ def _forward(
         loss_ctx: MoELossContextDict | None,
         return_router_logits: bool = False,
     ) -> MoEModelOutputs:
+        self._prepare_seq_ctx_topk_cache([seq_ctx])
         input_ids = seq_ctx.input_ids
         position_ids = seq_ctx.position_ids
 
@@ -854,7 +866,9 @@ def _forward(
                 input_ids=input_ids.clone() if input_ids is not None else None,
                 position_ids=position_ids.clone(),
                 inputs_embeds=seq_ctx.inputs_embeds.clone() if seq_ctx.inputs_embeds is not None else None,
-                dsa_topk_cache=DSATopKCacheState(),
+                dsa_topk_cache=DSATopKCacheState(
+                    offload_slot=seq_ctx.dsa_topk_cache.offload_slot,
+                ),
             )
             # MTP uses its own mask; main mask's non-pad indices do not apply.
             mtp_nonpad_indices = torch.nonzero(mtp_seq_ctx.mask, as_tuple=True)[1]

xtuner/v1/module/attention/dsa_topk_sharing.py

@@ -106,7 +106,7 @@ def after_recompute_release(self, seq_ctx: SequenceContext, source_layer_idx: in
         seq_ctx.dsa_topk_cache.indices.pop(source_layer_idx, None)
 
     def _offload_key(self, seq_ctx: SequenceContext, source_layer_idx: int) -> str:
-        return f"dsa_topk_{seq_ctx.dsa_topk_cache.context_id}_{source_layer_idx}"
+        return f"dsa_topk_{seq_ctx.dsa_topk_cache.offload_slot}_{source_layer_idx}"
 
 
 class ActivationOffloadedTopKResidency(GpuTopKResidency):
@@ -178,7 +178,16 @@ def after_original_forward_last_use(self, seq_ctx: SequenceContext, source_layer
             return
 
         key = self._offload_key(seq_ctx, source_layer_idx)
-        cpu_buffer = OffloadManager().get_or_create_pin_memory(key, topk_indices.shape, topk_indices.dtype)
+        # A slot/source pair owns both the runtime entry and its reusable pinned
+        # storage. Reusing it while still live would overwrite the earlier D2H
+        # result, so fail loudly if the scheduling contract is violated.
+        if OffloadManager().has_runtime_key(key):
+            raise RuntimeError(f"DSA top-k offload slot is still active: {key}")
+        cpu_buffer = OffloadManager().get_or_create_pin_memory(
+            key,
+            topk_indices.shape,
+            topk_indices.dtype,
+        )
         swap_tensor = SwapTensor(topk_indices, key, tensor_cpu=cpu_buffer)
         stream = self._stream_for_device(topk_indices.device)
         stream.wait_stream(torch.cuda.current_stream(topk_indices.device))

xtuner/v1/ops/sparse_mla/pytorch.py

@@ -68,7 +68,7 @@ def torch_dsa_topk_indices(
     topk = min(index_topk, kv_len)
     topk_scores, topk_indices = index_scores.topk(topk, dim=-1)
     topk_indices = topk_indices.masked_fill(topk_scores == -torch.inf, -1)
-    return topk_indices.squeeze(0).unsqueeze(1)
+    return topk_indices.squeeze(0).unsqueeze(1).to(torch.int32)
 
 
 def _packed_causal_mask(seq_ctx: SequenceContext, query_len: int, kv_len: int, device: torch.device) -> torch.Tensor:

xtuner/v1/ops/sparse_mla/tilelang.py

@@ -170,7 +170,7 @@ def _tilelang_dsa_topk_indices_from_ranges(
     topk = min(index_topk, k.shape[0])
     topk_scores, topk_indices = logits.topk(topk, dim=-1)
     topk_indices = topk_indices.masked_fill(topk_scores == -torch.inf, -1)
-    return topk_indices.to(torch.int64).unsqueeze(1)
+    return topk_indices.to(torch.int32).unsqueeze(1)
 
 
 @_tilelang_dsa_topk_indices_from_ranges.register_fake
@@ -183,7 +183,7 @@ def _(
     index_topk: int,
 ) -> Tensor:
     topk = min(index_topk, k.shape[0])
-    return torch.empty((q.shape[0], 1, topk), device=q.device, dtype=torch.int64)
+    return torch.empty((q.shape[0], 1, topk), device=q.device, dtype=torch.int32)
 
 
 @functools.cache

xtuner/v1/utils/activation_offload.py

@@ -229,6 +229,10 @@ def assert_exist(self, key):
     def exist(self, key):
         return key in self.items
 
+    def has_runtime_key(self, key):
+        """Return whether an offload key still owns a live runtime entry."""
+        return key in self.items or key in self.may_npu_tensors
+
     def assert_not_exist(self, key):
         if key not in self.items:
             raise RuntimeError(f"Key {key} already exist in items")

xtuner/v1/utils/misc.py

@@ -29,17 +29,34 @@
 
 def set_deterministic():
     os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"
-    # Inductor 会在 torch.compile 前读取 dynamic_scale_rblock；确定性模式必须尽早关闭。
-    # torch.use_deterministic_algorithms(True) 只会让 reduction 的初始候选收敛成一个 config；
-    # dynamic rblock 仍可能在 precompile 后追加一个缩小 R*_BLOCK 的 launcher，并由 runtime
-    # benchmark 在多个 launcher 中选择。不同 rank/run 一旦选到不同 reduction 分块，浮点累加
-    # 顺序就会变化，最终可能得到 bitwise 不同的梯度。
+    # Inductor 有两个独立的 reduction 调优入口：dynamic rblock 会在 precompile 后调整
+    # launcher，而 INNER reduction 会在 deterministic 检查前返回 persistent/contiguous
+    # 多个候选，由 runtime benchmark 选出配置。确定性模式必须同时关闭前者并固定后者。
+    # Triton compile subprocess 不会继承 Python monkey patch，因此也要在当前进程串行编译，
+    # 避免不同 rank/run 采用不同的浮点累加顺序。
     os.environ["TORCHINDUCTOR_DYNAMIC_SCALE_RBLOCK"] = "0"
+    os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
     from torch._inductor import config as inductor_config
+    from torch._inductor.runtime import triton_heuristics
 
     inductor_config.dynamic_scale_rblock = False
+    inductor_config.compile_threads = 1
     torch.use_deterministic_algorithms(True, warn_only=True)
 
+    original_reduction_configs = triton_heuristics._reduction_configs
+    if not getattr(original_reduction_configs, "_xtuner_deterministic_patched", False):
+
+        def deterministic_reduction_configs(*args, **kwargs):
+            return original_reduction_configs(*args, **kwargs)[:1]
+
+        setattr(deterministic_reduction_configs, "_xtuner_deterministic_patched", True)
+        setattr(
+            deterministic_reduction_configs,
+            "_xtuner_original_reduction_configs",
+            original_reduction_configs,
+        )
+        triton_heuristics._reduction_configs = deterministic_reduction_configs
+
 
 # https://github.com/python/cpython/issues/82300#issuecomment-2169035092
 if sys.version_info >= (3, 13):