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 = {
-                    # 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()
-                }
+                device_state = _get_offloaded_device_state(module, device)
                 output = functional_call(module, device_state, args=args, kwargs=kwargs)
                 module.forward = forward
                 return output