chore: import upstream snapshot with attribution

python/sglang/srt/mem_cache/dsa_cache_layer_split.py

@@ -0,0 +1,581 @@
+# Copyright 2023-2026 SGLang Team
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+
+"""Layer-sharded DSA KV cache pool for context-parallel prefill.
+
+``LayerSplitDSATokenToKVPool`` splits the DSA (DeepSeek Sparse Attention) GPU
+KV/indexer cache layers across context-parallel (CP) ranks so that each rank
+only materializes the layers it owns, reducing per-rank KV memory. When a rank
+needs to read a layer it does not own, the owning rank broadcasts that layer's
+buffer into a small per-rank remote scratch buffer.
+
+This subclass keeps the core ``KVCache`` / ``MLATokenToKVPool`` /
+``DSATokenToKVPool`` pools untouched: all sharding, broadcast, and remote-scratch
+bookkeeping lives here. Layer split is only ever enabled for DSA MLA models on
+PD prefill workers under prefill-CP (see
+``sglang.srt.layers.cp.utils.is_glm_dsa_cache_layer_split_enabled``).
+"""
+
+from __future__ import annotations
+
+import logging
+from contextlib import nullcontext
+from typing import TYPE_CHECKING, Optional
+
+import torch
+
+from sglang.srt.layers.attention.dsa import index_buf_accessor
+from sglang.srt.layers.cp.utils import get_layer_owner, get_layer_shard_range
+from sglang.srt.mem_cache.memory_pool import (
+    GPU_MEMORY_TYPE_KV_CACHE,
+    DSATokenToKVPool,
+    RadixAttention,
+    get_tensor_size_bytes,
+    maybe_detect_oob,
+    unwrap_write_loc,
+)
+from sglang.srt.runtime_context import get_parallel
+
+if TYPE_CHECKING:
+    from sglang.srt.managers.cache_controller import LayerDoneCounter
+
+logger = logging.getLogger(__name__)
+
+
+class LayerSplitDSATokenToKVPool(DSATokenToKVPool):
+    """DSA KV pool that shards layers across CP ranks with owner-broadcast reads."""
+
+    def __init__(
+        self,
+        *args,
+        layer_shard_rank: int,
+        layer_shard_size: int,
+        **kwargs,
+    ):
+        assert (
+            layer_shard_rank is not None and layer_shard_size > 1
+        ), "LayerSplitDSATokenToKVPool requires layer_shard_size > 1"
+        self.layer_shard_rank = layer_shard_rank
+        self.layer_shard_size = layer_shard_size
+        self.layer_shard_enabled = True
+        self.layer_broadcast_comm = None
+        super().__init__(*args, **kwargs)
+        # First global layer index owned by this rank (used by PD transfer to
+        # label the contiguous owned-buffer range).
+        my_start, _ = self._owned_local_layer_range()
+        self.layer_shard_start = self.start_layer + my_start
+
+    # ---- layer ownership helpers ------------------------------------------
+
+    def _local_layer_idx(self, layer_id: int) -> int:
+        return layer_id - self.start_layer
+
+    def _owned_local_layer_range(self) -> tuple[int, int]:
+        return get_layer_shard_range(
+            self.layer_shard_rank, self.layer_shard_size, self.layer_num
+        )
+
+    def _is_layer_owned(self, layer_id: int) -> bool:
+        local_idx = self._local_layer_idx(layer_id)
+        owned_start, owned_end = self._owned_local_layer_range()
+        return owned_start <= local_idx < owned_end
+
+    def _get_layer_owner_rank(self, layer_id: int) -> int:
+        return get_layer_owner(
+            self._local_layer_idx(layer_id), self.layer_shard_size, self.layer_num
+        )
+
+    def _log_layer_shard_plan(self) -> None:
+        partitions = []
+        for rank in range(self.layer_shard_size):
+            st, ed = get_layer_shard_range(rank, self.layer_shard_size, self.layer_num)
+            partitions.append(f"r{rank}:[{st},{ed})")
+        my_start, my_end = self._owned_local_layer_range()
+        logger.info(
+            "Layer shard plan (continuous): "
+            f"layer_num={self.layer_num}, shard_size={self.layer_shard_size}, "
+            f"rank={self.layer_shard_rank}, local=[{my_start},{my_end}), "
+            f"global=[{self.start_layer + my_start},{self.start_layer + my_end}), "
+            f"partitions={'; '.join(partitions)}"
+        )
+
+    # ---- broadcast plumbing -----------------------------------------------
+
+    def _init_layer_broadcast_comm(self) -> None:
+        cp_group = get_parallel().attn_cp_group
+        if cp_group.world_size <= 1 or cp_group.pynccl_comm is None:
+            return
+
+        from sglang.srt.distributed.device_communicators.pynccl import (
+            PyNcclCommunicator,
+        )
+
+        self.layer_broadcast_comm = PyNcclCommunicator(
+            group=cp_group.cpu_group,
+            device=cp_group.device,
+        )
+        logger.info(
+            "Initialized dedicated layer-shard broadcast NCCL communicator: "
+            f"rank={cp_group.rank_in_group}, world_size={cp_group.world_size}"
+        )
+
+    def _broadcast_tensor_from_owner(
+        self,
+        tensor: torch.Tensor,
+        layer_id: int,
+        src_tensor: Optional[torch.Tensor] = None,
+        use_layer_broadcast_comm: bool = False,
+    ) -> torch.Tensor:
+        owner_rank = self._get_layer_owner_rank(layer_id)
+        if self.layer_shard_rank == owner_rank:
+            assert src_tensor is not None
+            if tensor.data_ptr() != src_tensor.data_ptr():
+                tensor.copy_(src_tensor)
+
+        cp_group = get_parallel().attn_cp_group
+        comm = (
+            self.layer_broadcast_comm
+            if use_layer_broadcast_comm and self.layer_broadcast_comm is not None
+            else cp_group.pynccl_comm
+        )
+        if comm is not None:
+            # PyNcclCommunicator defaults to disabled=True (it is only enabled
+            # inside CUDA-graph capture via change_state). Without re-enabling it
+            # here, comm.broadcast() is a silent no-op and non-owner CP ranks read
+            # stale remote buffers, corrupting layer-split attention. Mirror the
+            # standard usage in parallel_state.py.
+            with comm.change_state(enable=True):
+                comm.broadcast(tensor, src=owner_rank)
+        else:
+            torch.distributed.broadcast(
+                tensor, src=owner_rank, group=cp_group.cpu_group
+            )
+        return tensor
+
+    # ---- buffer allocation (owned-only + remote scratch) ------------------
+
+    def _create_buffers(self):
+        self._log_layer_shard_plan()
+        with self.memory_saver_adapter.region(GPU_MEMORY_TYPE_KV_CACHE):
+            with (
+                torch.cuda.use_mem_pool(self.custom_mem_pool)
+                if self.custom_mem_pool
+                else nullcontext()
+            ):
+                # Owned layers get the full buffer; non-owned layers allocate a
+                # 0-row placeholder so ``kv_buffer`` stays index-aligned by layer.
+                self.kv_buffer = [
+                    torch.zeros(
+                        (
+                            (
+                                (self.size + self.page_size)
+                                if self._is_layer_owned(self.start_layer + i)
+                                else 0
+                            ),
+                            1,
+                            self.kv_cache_dim,
+                        ),
+                        dtype=self.store_dtype,
+                        device=self.device,
+                    )
+                    for i in range(self.layer_num)
+                ]
+                self.remote_kv_buffer = torch.empty(
+                    (self.size + self.page_size, 1, self.kv_cache_dim),
+                    dtype=self.store_dtype,
+                    device=self.device,
+                )
+                self.remote_kv_layer_id: Optional[int] = None
+                self.device_module = torch.get_device_module(self.device)
+                self.kv_broadcast_stream = self.device_module.Stream()
+                self.pending_remote_kv_layer_id: Optional[int] = None
+                self.pending_remote_kv_broadcast = False
+        self._init_layer_broadcast_comm()
+
+    def _create_index_buffers(self):
+        num_pages = (self.index_buf_size + self.page_size + 1) // self.page_size
+        with (
+            torch.cuda.use_mem_pool(self.custom_mem_pool)
+            if self.custom_mem_pool
+            else nullcontext()
+        ):
+            self.index_k_with_scale_buffer = [
+                torch.zeros(
+                    self._index_buffer_shape(
+                        num_pages if self._is_layer_owned(self.start_layer + i) else 0
+                    ),
+                    dtype=self.index_k_with_scale_buffer_dtype,
+                    device=self.device,
+                )
+                for i in range(self.layer_num)
+            ]
+            self.remote_index_k_with_scale_buffer = torch.empty(
+                self._index_buffer_shape(num_pages),
+                dtype=self.index_k_with_scale_buffer_dtype,
+                device=self.device,
+            )
+            self.remote_index_layer_id: Optional[int] = None
+
+    def _clear_buffers(self):
+        del self.kv_buffer
+        del self.remote_kv_buffer
+        del self.remote_index_k_with_scale_buffer
+        del self.index_k_with_scale_buffer
+
+    # ---- MLA latent KV: owned-only writes, owner-broadcast reads ----------
+
+    def get_kv_size_bytes(self):
+        kv_size_bytes = 0
+        for kv_cache in self.kv_buffer:
+            kv_size_bytes += get_tensor_size_bytes(kv_cache)
+        for index_k_cache in self.index_k_with_scale_buffer:
+            kv_size_bytes += get_tensor_size_bytes(index_k_cache)
+        return kv_size_bytes
+
+    def get_contiguous_buf_infos(self):
+        # Only report buffers owned by the current CP rank; non-owned layers
+        # are empty and are pulled from their owner via PD transfer.
+        owned_layer_ids = [
+            i
+            for i in range(self.layer_num)
+            if self._is_layer_owned(self.start_layer + i)
+        ]
+        kv_data_ptrs = [self.kv_buffer[i].data_ptr() for i in owned_layer_ids]
+        kv_data_lens = [self.kv_buffer[i].nbytes for i in owned_layer_ids]
+        kv_item_lens = [
+            self.kv_buffer[i][0].nbytes * self.page_size for i in owned_layer_ids
+        ]
+        return kv_data_ptrs, kv_data_lens, kv_item_lens
+
+    def get_key_buffer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+
+        kv_buffer = self._get_broadcastable_kv_buffer(layer_id)
+        if self.store_dtype != self.dtype:
+            return kv_buffer.view(self.dtype)
+        return kv_buffer
+
+    def get_value_buffer(self, layer_id: int):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+
+        kv_buffer = self._get_broadcastable_kv_buffer(layer_id)
+        if self.store_dtype != self.dtype:
+            return kv_buffer[..., : self.kv_lora_rank].view(self.dtype)
+        return kv_buffer[..., : self.kv_lora_rank]
+
+    def set_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc_info,
+        cache_k: torch.Tensor,
+        cache_v: torch.Tensor,
+    ):
+        loc, _, _ = unwrap_write_loc(loc_info)
+        maybe_detect_oob(loc, 0, self.size + self.page_size, "set_kv_buffer (MLA)")
+        layer_id = layer.layer_id
+        assert not self.dsa_kv_cache_store_fp8
+        # A write invalidates any cached remote copy for this layer.
+        if self.pending_remote_kv_layer_id == layer_id:
+            self._finalize_pending_kv_broadcast(set_remote_layer_id=False)
+        if self.remote_kv_layer_id == layer_id:
+            self.remote_kv_layer_id = None
+        if not self._is_layer_owned(layer_id):
+            return
+        if cache_k.dtype != self.dtype:
+            cache_k = cache_k.to(self.dtype)
+        if self.store_dtype != self.dtype:
+            self.kv_buffer[layer_id - self.start_layer][loc] = cache_k.view(
+                self.store_dtype
+            )
+        else:
+            self.kv_buffer[layer_id - self.start_layer][loc] = cache_k
+
+    def set_mla_kv_buffer(
+        self,
+        layer: RadixAttention,
+        loc: torch.Tensor,
+        cache_k_nope: torch.Tensor,
+        cache_k_rope: torch.Tensor,
+    ):
+        maybe_detect_oob(loc, 0, self.size + self.page_size, "set_mla_kv_buffer (MLA)")
+        layer_id = layer.layer_id
+        if self.pending_remote_kv_layer_id == layer_id:
+            self._finalize_pending_kv_broadcast(set_remote_layer_id=True)
+        remote_kv_updatable = self.remote_kv_layer_id == layer_id
+        if remote_kv_updatable:
+            self._write_mla_kv_buffer(
+                self.remote_kv_buffer, loc, cache_k_nope, cache_k_rope
+            )
+        if not self._is_layer_owned(layer_id):
+            return
+        self._write_mla_kv_buffer(
+            self.kv_buffer[layer_id - self.start_layer],
+            loc,
+            cache_k_nope,
+            cache_k_rope,
+        )
+        if not remote_kv_updatable and self.remote_kv_layer_id == layer_id:
+            self.remote_kv_layer_id = None
+
+    def _finalize_pending_kv_broadcast(
+        self, *, set_remote_layer_id: bool = True
+    ) -> None:
+        if not self.pending_remote_kv_broadcast:
+            return
+        self.device_module.current_stream().wait_stream(self.kv_broadcast_stream)
+        self.pending_remote_kv_broadcast = False
+        if set_remote_layer_id and self.pending_remote_kv_layer_id is not None:
+            self.remote_kv_layer_id = self.pending_remote_kv_layer_id
+        self.pending_remote_kv_layer_id = None
+
+    def prefetch_kv_buffer(
+        self,
+        layer_id: int,
+        layer_transfer_counter: Optional[LayerDoneCounter] = None,
+        layer_transfer_idx: Optional[int] = None,
+    ) -> None:
+        """Kick off an async owner-broadcast of ``layer_id``'s latent KV.
+
+        Called ahead of the layer's attention so the remote scratch buffer is
+        ready by the time a non-owner rank reads it (see the prefetch wiring in
+        ``DeepseekV2DecoderLayer``).
+        """
+        if self.remote_kv_layer_id == layer_id:
+            return
+        if self.pending_remote_kv_broadcast:
+            if self.pending_remote_kv_layer_id == layer_id:
+                return
+            self._finalize_pending_kv_broadcast(set_remote_layer_id=False)
+
+        local_idx = self._local_layer_idx(layer_id)
+        src_tensor = (
+            self.kv_buffer[local_idx] if self._is_layer_owned(layer_id) else None
+        )
+        if self.layer_broadcast_comm is None:
+            self._broadcast_tensor_from_owner(
+                self.remote_kv_buffer,
+                layer_id,
+                src_tensor=src_tensor,
+                use_layer_broadcast_comm=True,
+            )
+            self.remote_kv_layer_id = layer_id
+            return
+
+        self.kv_broadcast_stream.wait_stream(self.device_module.current_stream())
+        with self.device_module.stream(self.kv_broadcast_stream):
+            if layer_transfer_counter is not None and layer_transfer_idx is not None:
+                layer_transfer_counter.wait_until(layer_transfer_idx)
+            self._broadcast_tensor_from_owner(
+                self.remote_kv_buffer,
+                layer_id,
+                src_tensor=src_tensor,
+                use_layer_broadcast_comm=True,
+            )
+        self.pending_remote_kv_layer_id = layer_id
+        self.pending_remote_kv_broadcast = True
+
+    def _get_broadcastable_kv_buffer(self, layer_id: int) -> torch.Tensor:
+        if self.pending_remote_kv_broadcast:
+            self._finalize_pending_kv_broadcast(
+                set_remote_layer_id=self.pending_remote_kv_layer_id == layer_id
+            )
+        if self.remote_kv_layer_id != layer_id:
+            local_idx = self._local_layer_idx(layer_id)
+            src_tensor = (
+                self.kv_buffer[local_idx] if self._is_layer_owned(layer_id) else None
+            )
+            self._broadcast_tensor_from_owner(
+                self.remote_kv_buffer,
+                layer_id,
+                src_tensor=src_tensor,
+                use_layer_broadcast_comm=True,
+            )
+            self.remote_kv_layer_id = layer_id
+        return self.remote_kv_buffer
+
+    def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor):
+        size_limit = self.size + self.page_size
+        maybe_detect_oob(tgt_loc, 0, size_limit, "move_kv_cache tgt_loc")
+        maybe_detect_oob(src_loc, 0, size_limit, "move_kv_cache src_loc")
+        if tgt_loc.numel() == 0:
+            return
+        tgt_loc_flat = tgt_loc.view(-1).long()
+        src_loc_flat = src_loc.view(-1).long()
+        for kv_cache in self.kv_buffer:
+            if kv_cache.shape[0] == 0:
+                continue
+            kv_cache[tgt_loc_flat] = kv_cache[src_loc_flat]
+        for index_k in self.index_k_with_scale_buffer:
+            if index_k.shape[0] == 0:
+                continue
+            index_k[tgt_loc_flat] = index_k[src_loc_flat]
+
+    # ---- DSA indexer buffer: owned-only writes, owner-broadcast reads -----
+
+    def get_broadcastable_index_k_with_scale_buffer(
+        self, layer_id: int
+    ) -> torch.Tensor:
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+        return self._get_broadcastable_index_buffer(layer_id)
+
+    def get_index_k_continuous(self, layer_id, seq_len, page_indices):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+        buf = self._get_broadcastable_index_buffer(layer_id)
+        return index_buf_accessor.GetK.execute(
+            self, buf, seq_len=seq_len, page_indices=page_indices
+        )
+
+    def get_index_k_scale_continuous(self, layer_id, seq_len, page_indices):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+        buf = self._get_broadcastable_index_buffer(layer_id)
+        return index_buf_accessor.GetS.execute(
+            self, buf, seq_len=seq_len, page_indices=page_indices
+        )
+
+    def get_index_k_scale_buffer(
+        self, layer_id, seq_len_tensor, page_indices, seq_len_sum, max_seq_len
+    ):
+        if self.layer_transfer_counter is not None:
+            self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
+        buf = self._get_broadcastable_index_buffer(layer_id)
+        # Overlap the latent-KV owner-broadcast with the indexer read.
+        self.prefetch_kv_buffer(layer_id)
+        return index_buf_accessor.GetKAndS.execute(
+            self,
+            buf,
+            page_indices=page_indices,
+            seq_len_tensor=seq_len_tensor,
+            seq_len_sum=seq_len_sum,
+            max_seq_len=max_seq_len,
+        )
+
+    def set_index_k_scale_buffer(self, layer_id, loc, index_k, index_k_scale) -> None:
+        self.invalidate_index_buffer_for_layer(layer_id)
+        if not self._is_layer_owned(layer_id):
+            return
+        buf = self.index_k_with_scale_buffer[layer_id - self.start_layer]
+        index_buf_accessor.SetKAndS.execute(
+            pool=self, buf=buf, loc=loc, index_k=index_k, index_k_scale=index_k_scale
+        )
+
+    def invalidate_index_buffer_for_layer(self, layer_id: int) -> None:
+        if self.remote_index_layer_id == layer_id:
+            self.remote_index_layer_id = None
+
+    def _get_broadcastable_index_buffer(self, layer_id: int) -> torch.Tensor:
+        if self.remote_index_layer_id != layer_id:
+            local_idx = self._local_layer_idx(layer_id)
+            src_tensor = (
+                self.index_k_with_scale_buffer[local_idx]
+                if self._is_layer_owned(layer_id)
+                else None
+            )
+            self._broadcast_tensor_from_owner(
+                self.remote_index_k_with_scale_buffer,
+                layer_id,
+                src_tensor=src_tensor,
+            )
+            self.remote_index_layer_id = layer_id
+        return self.remote_index_k_with_scale_buffer
+
+    def get_state_buf_infos(self):
+        owned_layer_ids = [
+            i
+            for i in range(self.layer_num)
+            if self._is_layer_owned(self.start_layer + i)
+        ]
+        data_ptrs = [
+            self.index_k_with_scale_buffer[i].data_ptr() for i in owned_layer_ids
+        ]
+        data_lens = [self.index_k_with_scale_buffer[i].nbytes for i in owned_layer_ids]
+        item_lens = [
+            self.index_k_with_scale_buffer[i][0].nbytes for i in owned_layer_ids
+        ]
+        return data_ptrs, data_lens, item_lens
+
+    # ---- HiCache CPU offload: skip empty (non-owned) layers ---------------
+
+    def get_cpu_copy(self, indices, mamba_indices=None):
+        from sglang.srt.utils import current_platform
+
+        current_platform.synchronize()
+        kv_cache_cpu = []
+        chunk_size = self.cpu_offloading_chunk_size
+        for layer_id in range(self.layer_num):
+            kv_cache_cpu.append([])
+            if self.kv_buffer[layer_id].shape[0] == 0:
+                continue
+            for i in range(0, len(indices), chunk_size):
+                chunk_indices = indices[i : i + chunk_size]
+                kv_cpu = self.kv_buffer[layer_id][chunk_indices].to(
+                    "cpu", non_blocking=True
+                )
+                kv_cache_cpu[-1].append(kv_cpu)
+        current_platform.synchronize()
+
+        page_indices = indices[:: self.page_size] // self.page_size
+        torch.cuda.synchronize()
+        index_k_cpu = []
+        page_chunk_size = max(1, chunk_size // self.page_size)
+        for layer_id in range(self.layer_num):
+            index_k_cpu.append([])
+            if self.index_k_with_scale_buffer[layer_id].shape[0] == 0:
+                continue
+            for i in range(0, len(page_indices), page_chunk_size):
+                chunk_page_indices = page_indices[i : i + page_chunk_size]
+                idx_cpu = self.index_k_with_scale_buffer[layer_id][
+                    chunk_page_indices
+                ].to("cpu", non_blocking=True)
+                index_k_cpu[-1].append(idx_cpu)
+        torch.cuda.synchronize()
+        return {"kv": kv_cache_cpu, "index_k": index_k_cpu}
+
+    def load_cpu_copy(self, kv_cache_cpu_dict, indices, mamba_indices=None):
+        from sglang.srt.utils import current_platform
+
+        kv_cache_cpu = kv_cache_cpu_dict["kv"]
+        current_platform.synchronize()
+        chunk_size = self.cpu_offloading_chunk_size
+        for layer_id in range(self.layer_num):
+            if self.kv_buffer[layer_id].shape[0] == 0:
+                continue
+            for i in range(0, len(indices), chunk_size):
+                chunk_indices = indices[i : i + chunk_size]
+                kv_cpu = kv_cache_cpu[layer_id][i // chunk_size]
+                assert kv_cpu.shape[0] == len(chunk_indices)
+                kv_chunk = kv_cpu.to(self.kv_buffer[layer_id].device, non_blocking=True)
+                self.kv_buffer[layer_id][chunk_indices] = kv_chunk
+        current_platform.synchronize()
+
+        page_indices = indices[:: self.page_size] // self.page_size
+        index_k_cpu = kv_cache_cpu_dict["index_k"]
+        torch.cuda.synchronize()
+        page_chunk_size = max(1, chunk_size // self.page_size)
+        for layer_id in range(self.layer_num):
+            if self.index_k_with_scale_buffer[layer_id].shape[0] == 0:
+                continue
+            for i in range(0, len(page_indices), page_chunk_size):
+                chunk_page_indices = page_indices[i : i + page_chunk_size]
+                idx_cpu = index_k_cpu[layer_id][i // page_chunk_size]
+                assert idx_cpu.shape[0] == len(chunk_page_indices)
+                idx_chunk = idx_cpu.to(
+                    self.index_k_with_scale_buffer[layer_id].device, non_blocking=True
+                )
+                self.index_k_with_scale_buffer[layer_id][chunk_page_indices] = idx_chunk
+        torch.cuda.synchronize()