Files
sglang/test/registered/unit
laoyao0822 f8fca72635 Expand prefill CP KV capacity by sharding persistent NSA KV
Prefill CP previously replicated NSA/MLA persistent KV on every CP rank, so CP8 consumed eight copies of KV memory while exposing only one rank of logical cache capacity. This change splits logical KV locs from per-rank physical storage, shards MLA latent KV and NSA index K/scale by deterministic page ownership, and keeps existing NSA attention kernels working through a full-view runtime materialization layer.

Mooncake PD transfer now sends each prefill CP rank's owned physical pages with explicit logical page positions so non-CP decode can reconstruct full-layout KV. The implementation is guarded by an explicit server flag and startup checks, and the design documentation records the implemented scope, debug environment, and Phase 3 boundary.

Constraint: Phase 2 must preserve existing NSA attention/index kernels via runtime full-view materialization
Constraint: Decode side remains non-CP and receives full KV through Mooncake
Rejected: Shard-aware NSA attention in this change | belongs to Phase 3 because it requires distributed topk/softmax/output contracts
Rejected: Request-contiguous CP ownership | unstable under chunked prefill and tied to attention split mode
Confidence: medium
Scope-risk: broad
Directive: Do not enable round-robin CP shared KV without wiring runtime materialization/PD transfer contracts for that split mode
Directive: Keep SGLANG_DEBUG_CP_SHARED_KV disabled for perf measurements; it intentionally enables CUDA-syncing diagnostics
Tested: Remote py_compile for shared-KV touched Python files in g0034 container
Tested: Remote pytest selected cp_shared/shared_kv/nsa suite: 37 passed, 34 deselected
Not-tested: Full GLM5 multi-node throughput/regression run after final doc update
Not-tested: Phase 3 shard-aware runtime, round-robin CP mode, and non-Mooncake PD backends
2026-04-26 04:11:17 +08:00
..

Unit Tests

Component-level tests that do not launch a server or load model weights. Tests can use CPU or GPU — the key criterion is no server process.

Quick Start

  1. Find the source file under python/sglang/srt/.
  2. Create the corresponding test here, mirroring the source tree:
    srt/mem_cache/radix_cache.py       →  unit/mem_cache/test_radix_cache.py
    srt/sampling/sampling_params.py    →  unit/sampling/test_sampling_params.py
    
  3. Register for CI at the top of the file (after imports, before test classes):
    from sglang.test.ci.ci_register import register_cpu_ci
    register_cpu_ci(est_time=5, suite="stage-a-test-cpu")
    # or: register_cuda_ci(est_time=10, suite="stage-b-test-1-gpu-small")
    
  4. Run locally:
    pytest test/registered/unit/ -v            # all unit tests
    pytest test/registered/unit/mem_cache/ -v  # one module
    
  5. Run with coverage:
    # summary
    pytest test/registered/unit/ --cov --cov-config=.coveragerc -v
    
    # PR incremental check (require ≥60% on changed lines)
    pytest test/registered/unit/ --cov --cov-config=.coveragerc --cov-report=xml
    diff-cover coverage.xml --compare-branch=origin/main --fail-under=60
    

Example

"""Unit tests for <module> — no server, no model loading."""

from sglang.test.ci.ci_register import register_cpu_ci

register_cpu_ci(est_time=5, suite="stage-a-test-cpu")

import unittest

from sglang.srt.<module> import TargetClass
from sglang.test.test_utils import CustomTestCase


class TestTargetClass(CustomTestCase):
    def test_basic_behavior(self):
        obj = TargetClass(...)
        self.assertEqual(obj.method(), expected)


if __name__ == "__main__":
    unittest.main()

Rules

  • No popen_launch_server() or Engine(...).
  • No model weight loading.
  • Use CustomTestCase (from sglang.test.test_utils, adds CI retry).
  • Use unittest.mock for dependencies that are expensive to construct.