Files
sglang/test/registered/unit
laoyao0822 4414db594c Reduce spec-v2 overlap copy overhead
Spec-v2 decode was paying unnecessary scheduler-side copy and replay costs: hidden states were copied back to CPU even when not requested, overlap result D2H ran on the forward stream, and draft graph replay issued several small copies separately.

Port the focused upstream optimizations without taking the larger spec-v2 relay refactor: gate hidden-state D2H on return_hidden_states, run result copies on copy_stream with pinned async D2H, use stride arange for draft select_index, and group small draft graph replay copies while keeping the large hidden-state DMA copy separate.

Constraint: Current branch carries local GLM/NSA/PD/CP changes, so upstream spec-v2 large refactors are not safe to cherry-pick wholesale.

Rejected: Cherry-pick upstream spec-v2 relay/dataclass refactor | too broad for this performance fix and conflicts with current branch structure

Rejected: Copy hidden states unconditionally | default chat output does not consume them after draft extend

Confidence: medium

Scope-risk: moderate

Directive: Keep result-copy lifetime tied to copy_done; do not move copy_to_cpu back onto forward_stream without measuring decode throughput.

Tested: local py_compile for changed production files and new test

Tested: local git diff --check

Tested: remote g0034 cjy-glm5-new as ubuntu on /sgl-workspace/sglang-tai: pytest -q -p no:cacheprovider test_generation_batch_result_copy.py test_eagle_worker_v2_cp_hidden.py test_spec_utils.py

Not-tested: full two-node decode throughput A/B after restart
2026-06-29 03:16:14 +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.