diff --git a/python/sglang/srt/environ.py b/python/sglang/srt/environ.py index 47b6c8f6e..dafeac92e 100644 --- a/python/sglang/srt/environ.py +++ b/python/sglang/srt/environ.py @@ -220,6 +220,11 @@ class Envs: # large bs) but coarser overlap. 1 = per-layer. SGLANG_CP_SHARED_KV_PER_LAYER_GROUP = EnvInt(8) SGLANG_CP_SHARED_KV_USE_TAI_MATERIALIZE = EnvBool(False) + # NSA MQA logits are materialized as fp32 [q, k] buffers inside DeepGEMM. + # Lower values split query rows more aggressively to cap peak temporary memory. + SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION = EnvFloat(0.2) + # Optional hard cap for rows per MQA-logits chunk. 0 = use memory budget. + SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS = EnvInt(0) SGLANG_CP_SHARED_KV_FUSED_MLA_STORE = EnvBool(False) SGLANG_CP_SHARED_KV_FUSED_INDEX_MQA_PREPARE = EnvBool(False) SGLANG_CP_SHARED_KV_ENABLE_MLA_PREFETCH = EnvBool(False) diff --git a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py index 537c3b4cc..10f0cfd69 100644 --- a/python/sglang/srt/layers/attention/nsa/nsa_indexer.py +++ b/python/sglang/srt/layers/attention/nsa/nsa_indexer.py @@ -418,10 +418,13 @@ class Indexer(MultiPlatformOp): # torch.cuda.mem_get_info host sync on the prefill critical path. (upstream PR #25299) _MQA_LOGITS_BYTES_PER_ELEM = 4 _MQA_LOGITS_STATIC_SKIP_ELEMS = 8_000_000 - _MQA_LOGITS_FREE_MEM_FRACTION = 0.5 _MQA_LOGITS_TOTAL_MEM_FRACTION = 0.3 _mqa_logits_budget_bytes: Dict[int, int] = {} + @staticmethod + def _mqa_logits_free_mem_fraction() -> float: + return envs.SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION.get() + def __init__( self, hidden_size: int, @@ -1124,6 +1127,7 @@ class Indexer(MultiPlatformOp): return topk_result def _get_mqa_logits_budget_bytes(self, device_index: int) -> int: + free_mem_fraction = self._mqa_logits_free_mem_fraction() # Cache the MQA-logits byte budget per device. torch.cuda.mem_get_info # host-syncs, so query free memory at most once (after the first real # prefill) and cap it by the workload-independent serving-memory headroom @@ -1140,7 +1144,7 @@ class Indexer(MultiPlatformOp): else: static_free_mem = int(total_mem * max(0.0, 1.0 - mem_fraction_static)) static_budget = min( - int(static_free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), + int(static_free_mem * free_mem_fraction), total_mem_budget, ) static_budget = max(1, static_budget) @@ -1152,7 +1156,7 @@ class Indexer(MultiPlatformOp): free_mem, _ = torch.cuda.mem_get_info(device_index) budget_bytes = min( - int(free_mem * self._MQA_LOGITS_FREE_MEM_FRACTION), static_budget + int(free_mem * free_mem_fraction), static_budget ) budget_bytes = max(1, budget_bytes) self._mqa_logits_budget_bytes[device_index] = budget_bytes @@ -1175,6 +1179,98 @@ class Indexer(MultiPlatformOp): need_chunk = logits_bytes > logits_budget_bytes return need_chunk, logits_budget_bytes + def _mqa_logits_chunk_max_rows( + self, num_q: int, num_k: int, logits_budget_bytes: int + ) -> int: + explicit_max_rows = int(envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.get()) + if explicit_max_rows > 0: + return min(max(1, explicit_max_rows), max(1, num_q)) + + bytes_per_row = num_k * self._MQA_LOGITS_BYTES_PER_ELEM + max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) + return min(max_rows, max(1, num_q)) + + def _mqa_logits_topk_ragged_chunked( + self, + metadata: BaseIndexerMetadata, + q_fp8: torch.Tensor, + kv_fp8: Tuple[torch.Tensor, torch.Tensor], + weights: torch.Tensor, + ks: torch.Tensor, + ke: torch.Tensor, + *, + ke_offset: torch.Tensor, + topk_indices_offset_override: torch.Tensor, + forward_batch: Optional[ForwardBatch] = None, + ) -> torch.Tensor: + q_offset = int(q_fp8.shape[0]) + if q_offset == 0: + return torch.full( + (0, self.index_topk), + -1, + dtype=torch.int32, + device=q_fp8.device, + ) + + device_index = q_fp8.device.index + assert device_index is not None, "q_fp8 must be on an indexed CUDA device" + k_offset = int(kv_fp8[0].shape[0]) + need_chunk, logits_budget_bytes = self._should_chunk_mqa_logits( + q_offset, k_offset, device_index + ) + max_rows = self._mqa_logits_chunk_max_rows( + q_offset, k_offset, logits_budget_bytes + ) + need_chunk = need_chunk or max_rows < q_offset + + def _topk_for_slice(start: int, end: int) -> torch.Tensor: + with self._with_real_sm_count(): + if _is_hip: + from aiter.ops.triton.fp8_mqa_logits import fp8_mqa_logits + + kv, scale = kv_fp8 + logits = fp8_mqa_logits( + q_fp8[start:end], + kv, + scale, + weights[start:end], + ks[start:end], + ke[start:end], + ) + else: + logits = deep_gemm.fp8_mqa_logits( + q_fp8[start:end], + kv_fp8, + weights[start:end], + ks[start:end], + ke[start:end], + clean_logits=False, + ) + return metadata.topk_transform( + logits, + self.index_topk, + ks=ks[start:end], + ke_offset=ke_offset[start:end], + topk_indices_offset_override=topk_indices_offset_override[start:end], + ) + + if not need_chunk: + return _topk_for_slice(0, q_offset) + + topk_result = None + start = 0 + while start < q_offset: + end = min(start + max_rows, q_offset) + topk_chunk = _topk_for_slice(start, end) + if topk_result is None: + topk_result = topk_chunk.new_full( + (q_offset, topk_chunk.shape[1]), -1 + ) + topk_result[start:end] = topk_chunk + start = end + assert topk_result is not None + return topk_result + def _get_topk_ragged( self, enable_dual_stream: bool, @@ -1286,9 +1382,9 @@ class Indexer(MultiPlatformOp): return topk_result # Chunk path - bytes_per_row = k_offset * self._MQA_LOGITS_BYTES_PER_ELEM - max_rows = max(1, int(logits_budget_bytes // max(bytes_per_row, 1))) - max_rows = min(max_rows, q_offset) + max_rows = self._mqa_logits_chunk_max_rows( + q_offset, k_offset, logits_budget_bytes + ) global_topk_offset = metadata.attn_metadata.topk_indices_offset @@ -1588,23 +1684,16 @@ class Indexer(MultiPlatformOp): q_lens_list, dtype=torch.int32, device=q_fp8.device ) ke = ks + ke_offset - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, + topk_result = self._mqa_logits_topk_ragged_chunked( + metadata, + q_fp8, + kv_fp8, + weights, ks=ks, - cu_seqlens_q=actual_seq_q, + ke=ke, ke_offset=ke_offset, - batch_idx_list=batch_idx_list, topk_indices_offset_override=topk_indices_offset_override, + forward_batch=forward_batch, ) return topk_result else: @@ -1653,23 +1742,16 @@ class Indexer(MultiPlatformOp): q_lens_list, dtype=torch.int32, device=q_fp8.device ) ke = ks + ke_offset - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, + topk_result = self._mqa_logits_topk_ragged_chunked( + metadata, + q_fp8, + kv_fp8, + weights, ks=ks, - cu_seqlens_q=actual_seq_q, + ke=ke, ke_offset=ke_offset, - batch_idx_list=batch_idx_list, topk_indices_offset_override=topk_indices_offset_override, + forward_batch=forward_batch, ) return topk_result @@ -1748,23 +1830,16 @@ class Indexer(MultiPlatformOp): ke_offset = torch.cat(ke_offset_list, dim=0) ke = ks + ke_offset actual_seq_q = torch.cat(actual_seq_q_list, dim=0) - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) - topk_result = metadata.topk_transform( - logits, - self.index_topk, + topk_result = self._mqa_logits_topk_ragged_chunked( + metadata, + q_fp8, + kv_fp8, + weights, ks=ks, - cu_seqlens_q=actual_seq_q, + ke=ke, ke_offset=ke_offset, - batch_idx_list=batch_idx_list, topk_indices_offset_override=topk_indices_offset_override, + forward_batch=forward_batch, ) else: seq_len = int(forward_batch.seq_lens_cpu[batch_idx].item()) @@ -1863,16 +1938,6 @@ class Indexer(MultiPlatformOp): k_scale = k_scale[:kv_len].view(torch.float32).squeeze(-1).contiguous() kv_fp8 = (k_fp8, k_scale) ke = ke_offset - - with self._with_real_sm_count(): - logits = deep_gemm.fp8_mqa_logits( - q_fp8, - kv_fp8, - weights, - ks, - ke, - clean_logits=False, - ) topk_indices_offset_override = None cu_seqlens_q_topk_override = None if ( @@ -1884,7 +1949,17 @@ class Indexer(MultiPlatformOp): # produced a zero offset per query. Reuse `ks` and avoid the # post-MQA metadata kernels entirely. topk_indices_offset_override = ks - actual_seq_q_tensor = None + valid_topk_result = self._mqa_logits_topk_ragged_chunked( + metadata, + q_fp8, + kv_fp8, + weights, + ks=ks, + ke=ke, + ke_offset=ke_offset, + topk_indices_offset_override=topk_indices_offset_override, + forward_batch=forward_batch, + ) elif valid_q_count == actual_seq_q and actual_seq_q_cu_tensor is not None: cu_seqlens_q_topk_override = actual_seq_q_cu_tensor elif actual_seq_q_tensor is None or valid_q_count != actual_seq_q: @@ -1894,15 +1969,25 @@ class Indexer(MultiPlatformOp): cu_seqlens_q_topk_override[1] = actual_seq_q_tensor.reshape(-1)[0] elif actual_seq_q_tensor.ndim == 0: actual_seq_q_tensor = actual_seq_q_tensor.reshape(1) - valid_topk_result = metadata.topk_transform( - logits, - self.index_topk, - ks=ks, - cu_seqlens_q=actual_seq_q_tensor, - ke_offset=ke_offset, - topk_indices_offset_override=topk_indices_offset_override, - cu_seqlens_q_topk_override=cu_seqlens_q_topk_override, - ) + if topk_indices_offset_override is None: + with self._with_real_sm_count(): + logits = deep_gemm.fp8_mqa_logits( + q_fp8, + kv_fp8, + weights, + ks, + ke, + clean_logits=False, + ) + valid_topk_result = metadata.topk_transform( + logits, + self.index_topk, + ks=ks, + cu_seqlens_q=actual_seq_q_tensor, + ke_offset=ke_offset, + topk_indices_offset_override=topk_indices_offset_override, + cu_seqlens_q_topk_override=cu_seqlens_q_topk_override, + ) if valid_q_count == actual_seq_q: topk_result = valid_topk_result else: diff --git a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py index fc410405b..2805ae53b 100644 --- a/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py +++ b/test/registered/unit/mem_cache/test_cp_shared_kv_runtime.py @@ -5817,6 +5817,43 @@ class TestCpSharedKVTaiMaterializeIntegration(unittest.TestCase): self.assertIs(dense_pages, fallback_pages) logger.warning.assert_not_called() + def test_nsa_mqa_logits_chunk_budget_uses_env_fraction(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa import nsa_indexer + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + indexer._mqa_logits_budget_bytes = {} + + with envs.SGLANG_NSA_MQA_LOGITS_FREE_MEM_FRACTION.override(0.25), patch( + "sglang.srt.layers.attention.nsa.nsa_indexer.get_is_capture_mode", + return_value=False, + ), patch( + "sglang.srt.layers.attention.nsa.nsa_indexer.get_global_server_args", + return_value=SimpleNamespace(mem_fraction_static=0.5), + ), patch.object( + nsa_indexer.torch.cuda, "get_device_properties" + ) as props, patch.object( + nsa_indexer.torch.cuda, "mem_get_info", return_value=(80_000, 100_000) + ): + props.return_value = SimpleNamespace(total_memory=100_000) + self.assertEqual(indexer._get_mqa_logits_budget_bytes(0), 12_500) + + def test_nsa_mqa_logits_chunk_max_rows_overrides_budget_rows(self): + from sglang.srt.environ import envs + from sglang.srt.layers.attention.nsa.nsa_indexer import Indexer + + indexer = object.__new__(Indexer) + with envs.SGLANG_NSA_MQA_LOGITS_CHUNK_MAX_ROWS.override(128): + self.assertEqual( + indexer._mqa_logits_chunk_max_rows( + num_q=1024, + num_k=4096, + logits_budget_bytes=4096 * 4 * 512, + ), + 128, + ) + if __name__ == "__main__": unittest.main()