diff --git a/python/sglang/kernels/ops/attention/dsa/transform_index.py b/python/sglang/kernels/ops/attention/dsa/transform_index.py index e0aabe9f8..ae5abf381 100644 --- a/python/sglang/kernels/ops/attention/dsa/transform_index.py +++ b/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: diff --git a/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py b/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py index c36f892c5..7eb1ea012 100644 --- a/python/sglang/srt/layers/attention/dsa/dsa_indexer_kpool.py +++ b/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 ): diff --git a/python/sglang/srt/utils/offloader.py b/python/sglang/srt/utils/offloader.py index 2170ce9c5..acee05d38 100644 --- a/python/sglang/srt/utils/offloader.py +++ b/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 diff --git a/test/registered/kernels/ops/attention/test_dsa_transform_index.py b/test/registered/kernels/ops/attention/test_dsa_transform_index.py index 82910f6f4..4598bbe38 100644 --- a/test/registered/kernels/ops/attention/test_dsa_transform_index.py +++ b/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,