Support GLM-5.3-Flash hybrid attention CPU offload and PD index mapping (#40310)

Co-authored-by: Xinyuan Tong <115166877+JustinTong0323@users.noreply.github.com>

python/sglang/kernels/ops/attention/dsa/transform_index.py

@@ -220,7 +220,8 @@ def transform_index_page_table_prefill_fast(
     cu_seqlens_q: Optional[torch.Tensor] = None,
 ) -> torch.Tensor:
     assert page_size == 1
-    assert topk_indices.shape[1] == 2048
+    assert topk_indices.ndim == 2
+    assert topk_indices.shape[1] > 0
     real_num_tokens = sum(extend_lens_cpu)
     result = _allocate_prefill_result(topk_indices, real_num_tokens, output_num_tokens)
     if real_num_tokens == 0:

python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py

@@ -52,6 +52,12 @@ if TYPE_CHECKING:
     from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
 
 
+def _should_fuse_kpool_topk(metadata: BaseIndexerMetadata) -> bool:
+    return envs.SGLANG_DSA_FUSE_TOPK.get() and not getattr(
+        metadata, "force_unfused_topk", False
+    )
+
+
 class IndexerKPool(MultiPlatformOp):
     def __init__(
         self,
@@ -784,7 +790,7 @@ class IndexerKPool(MultiPlatformOp):
         paged_page_table: Optional[torch.Tensor] = None,
         paged_page_table_row_index: Optional[torch.Tensor] = None,
     ) -> Tuple[Optional[torch.Tensor], Optional[torch.Tensor], Optional[torch.Tensor]]:
-        if not envs.SGLANG_DSA_FUSE_TOPK.get():
+        if not _should_fuse_kpool_topk(metadata):
             return None, None, None
 
         topk_method = metadata.topk_transform_method
@@ -977,7 +983,7 @@ class IndexerKPool(MultiPlatformOp):
         page_table_all = None
         page_table_row_index_all = None
         topk_offsets_all = None
-        if envs.SGLANG_DSA_FUSE_TOPK.get():
+        if _should_fuse_kpool_topk(metadata):
             if topk_method == TopkTransformMethod.PAGED:
                 page_table_all = plan.ragged_paged_page_table
                 page_table_row_index_all = plan.ragged_paged_page_table_row_index
@@ -1209,7 +1215,7 @@ class IndexerKPool(MultiPlatformOp):
             page_table_local = None
             topk_offsets_local = None
             if (
-                envs.SGLANG_DSA_FUSE_TOPK.get()
+                _should_fuse_kpool_topk(metadata)
                 and topk_method == TopkTransformMethod.PAGED
             ):
                 page_table_local = (
@@ -1219,7 +1225,7 @@ class IndexerKPool(MultiPlatformOp):
                 )
                 page_table_local = page_table_local.unsqueeze(0).expand(q_len, -1)
             elif (
-                envs.SGLANG_DSA_FUSE_TOPK.get()
+                _should_fuse_kpool_topk(metadata)
                 and topk_method == TopkTransformMethod.RAGGED
                 and topk_offsets is not None
             ):

python/sglang/srt/utils/offloader.py

@@ -85,6 +85,57 @@ def create_offloader(dp_rank: int):
     return NoopOffloader()
 
 
+def _get_offloaded_device_state(module: torch.nn.Module, device: torch.device):
+    transferred = {}
+    device_state = {}
+    for name, value in module.state_dict(keep_vars=True).items():
+        key = id(value)
+        if key not in transferred:
+            transferred[key] = value.detach().to(device, non_blocking=True)
+        device_state[name] = transferred[key]
+    return device_state
+
+
+def _get_resident_parameter_ids(module: torch.nn.Module):
+    # functional_call only replaces registered parameters and buffers, so cached
+    # tensors held as ordinary attributes need their backing weights to stay put.
+    resident = set()
+    for owner in module.modules():
+        # MLA post_load_weights derives w_kc/w_vc from kv_b_proj weights on
+        # their current device. These attributes already exist before loading.
+        projection = getattr(owner, "kv_b_proj", None)
+        if (
+            isinstance(projection, torch.nn.Module)
+            and hasattr(owner, "w_kc")
+            and hasattr(owner, "w_vc")
+        ):
+            resident.update(id(parameter) for parameter in projection.parameters())
+        # KDA caches a storage-sharing view of qkv_conv1d.weight in conv_weights
+        # during construction. Offloading the weight would leave that view stale.
+        projection = getattr(owner, "qkv_conv1d", None)
+        attention = getattr(owner, "attn", None)
+        if isinstance(projection, torch.nn.Module) and isinstance(
+            attention, torch.nn.Module
+        ):
+            weight = getattr(projection, "weight", None)
+            cached = getattr(attention, "conv_weights", None)
+            if (
+                isinstance(weight, torch.nn.Parameter)
+                and isinstance(cached, torch.Tensor)
+                and cached.device == weight.device
+            ):
+                weight_storage = weight.untyped_storage()
+                cached_storage = cached.untyped_storage()
+                if (
+                    weight_storage.nbytes() > 0
+                    and weight_storage.data_ptr() != 0
+                    and weight_storage.nbytes() == cached_storage.nbytes()
+                    and weight_storage.data_ptr() == cached_storage.data_ptr()
+                ):
+                    resident.add(id(weight))
+    return resident
+
+
 class OffloaderV1(BaseOffloader):
     def __init__(self, cpu_offload_max_bytes: int):
         self._cpu_offload_bytes = 0
@@ -114,7 +165,10 @@ class OffloaderV1(BaseOffloader):
         # offload parameters to CPU
         # use pin_memory if possible, which helps cudagraph capture speed
         offloaded_parameters = False
+        resident_parameter_ids = _get_resident_parameter_ids(module)
         for p in module.parameters():
+            if id(p) in resident_parameter_ids:
+                continue
             if self._cpu_offload_bytes >= self._cpu_offload_max_bytes:
                 # we use per-parameter offloading
                 # one module might have some parameters offloaded and some not
@@ -139,12 +193,7 @@ class OffloaderV1(BaseOffloader):
 
             def forward(*args, **kwargs):
                 module.forward = original_forward
-                device_state = {
+                device_state = _get_offloaded_device_state(module, device)
-                    # here we blindly call `to(device)`
-                    # if the parameter is already on the device, it will be a no-op
-                    k: v.to(device, non_blocking=True)
-                    for k, v in module.state_dict().items()
-                }
                 output = functional_call(module, device_state, args=args, kwargs=kwargs)
                 module.forward = forward
                 return output

test/registered/kernels/ops/attention/test_dsa_transform_index.py

@@ -7,6 +7,7 @@ import sglang.kernels.ops.attention.dsa.transform_index as transform_index_modul
 from sglang.kernels.ops.attention.dsa.transform_index import (
     transform_index_page_table_decode_fast,
     transform_index_page_table_prefill_fast,
+    transform_index_page_table_prefill_ref,
 )
 from sglang.test.ci.ci_register import register_cuda_ci
 from sglang.test.test_utils import CustomTestCase
@@ -235,6 +236,43 @@ class TestDSATransformIndex(CustomTestCase):
             output_padding=8,
         )
 
+    def test_prefill_kpool_tail_width(self):
+        extend_lens_cpu = [0, 3, 1, 0, 4]
+        real_num_tokens = sum(extend_lens_cpu)
+        topk_num_tokens = real_num_tokens + 5
+        output_num_tokens = topk_num_tokens + 7
+        context_length = 8192
+
+        # KPool can append three tail positions after the 2048 history indices.
+        tail_indices = torch.tensor(
+            [context_length - 3, context_length - 2, context_length - 1],
+            dtype=torch.int64,
+            device=self.device,
+        ).repeat(topk_num_tokens, 1)
+        topk_indices = torch.cat(
+            [self._make_topk(topk_num_tokens, context_length), tail_indices], dim=1
+        )
+        topk_indices[0, -1] = -1
+        self.assertEqual(topk_indices.shape[1], 2051)
+
+        for page_table_is_expanded in (False, True):
+            with self.subTest(page_table_is_expanded=page_table_is_expanded):
+                page_table_rows = (
+                    real_num_tokens if page_table_is_expanded else len(extend_lens_cpu)
+                )
+                page_table = self._make_page_table(page_table_rows, context_length)
+                kwargs = dict(
+                    page_table=page_table,
+                    topk_indices=topk_indices,
+                    extend_lens_cpu=extend_lens_cpu,
+                    output_num_tokens=output_num_tokens,
+                    page_table_is_expanded=page_table_is_expanded,
+                )
+                expected = transform_index_page_table_prefill_ref(**kwargs)
+                actual = transform_index_page_table_prefill_fast(**kwargs)
+                torch.cuda.synchronize()
+                torch.testing.assert_close(actual, expected, rtol=0, atol=0)
+
     def test_large_batch_size(self):
         self._check_case(
             [1] * 8192,