diff --git a/tests/router/common.py b/tests/router/common.py index 99aa928767a6..c1856a8dc35d 100644 --- a/tests/router/common.py +++ b/tests/router/common.py @@ -2409,6 +2409,7 @@ def _test_router_decisions_disagg( request_plane: str = "nats", router_ais_config: Optional[dict[str, Any]] = None, enable_bootstrap: bool = False, + require_kv_hit: bool = False, ): """Validate KV cache prefix reuse in disaggregated prefill-decode setup via HTTP frontend. @@ -2431,6 +2432,10 @@ def _test_router_decisions_disagg( test_payload: Base test payload to send to /v1/chat/completions store_backend: Storage backend to use ("etcd" or "file"). Defaults to "etcd". router_ais_config: Optional AIS router perf-model config for frontend KV routing. + require_kv_hit: If True, also assert the router predicted a KV cache hit + (nvext.timing.kv_hit_rate > 0) for requests 2-4, which share a prefix + with the previous request. This catches worker/router block-hash + mismatches that routing to the same worker alone would not. Raises: AssertionError: If prefill_worker_ids differ across requests (prefix reuse failure) @@ -2542,6 +2547,13 @@ async def send_progressive_requests(): timing_info is not None ), f"Request {i + 1}: Expected timing info in final chunk, got None" verify_response_timing(timing_info, disagg=not enable_bootstrap) + if require_kv_hit and i > 0: + kv_hit_rate = timing_info.get("kv_hit_rate") + assert kv_hit_rate is not None and kv_hit_rate > 0, ( + f"Request {i + 1}: expected kv_hit_rate > 0 on a shared " + f"prefix, got {kv_hit_rate}; worker KV event hashes " + f"may not match the router's" + ) # Small delay between requests await asyncio.sleep(1) diff --git a/tests/router/e2e_harness.py b/tests/router/e2e_harness.py index a1ba23b1dc04..e9fd68bc8fee 100644 --- a/tests/router/e2e_harness.py +++ b/tests/router/e2e_harness.py @@ -423,7 +423,14 @@ def run_disagg_router_decisions_test( | None = None, test_payload: dict[str, Any] | None = None, test_kwargs: dict[str, Any] | None = None, + post_check: Callable[[Any, Any], None] | None = None, ): + """Run the disagg router decisions test against prefill and decode workers. + + ``post_check(prefill_workers, decode_workers)``, if set, runs after the + routing assertions while the workers are still up (e.g. to scrape their + /metrics). + """ shared_namespace = f"test-namespace-{generate_random_suffix()}" frontend_port = allocate_frontend_ports(request, 1)[0] @@ -447,6 +454,8 @@ def run_test(prefill_workers, decode_workers): request_plane=request_plane, **(test_kwargs or {}), ) + if post_check is not None: + post_check(prefill_workers, decode_workers) if worker_context_factory is not None: with worker_context_factory(shared_namespace) as workers: diff --git a/tests/router/test_router_e2e_with_sglang.py b/tests/router/test_router_e2e_with_sglang.py index 3d8c1dc46f70..f8dc29a08061 100644 --- a/tests/router/test_router_e2e_with_sglang.py +++ b/tests/router/test_router_e2e_with_sglang.py @@ -10,9 +10,11 @@ from typing import Any, Dict, Optional import pytest +import requests from tests.router.e2e_harness import ( ManagedEngineProcessMixin, + build_test_payload, run_basic_router_test, run_cache_salt_isolation_test, run_disagg_router_decisions_test, @@ -23,6 +25,7 @@ from tests.utils.constants import DynamoPortRange from tests.utils.gpu_args import build_gpu_mem_args from tests.utils.managed_process import ManagedProcess, check_health_ready +from tests.utils.payloads import SGLangSpecDecodeMetricsPayload from tests.utils.port_utils import ( allocate_contiguous_ports, allocate_port, @@ -95,6 +98,7 @@ def __init__( (see tests/README.md "SGLang KV tokens") - context_length: Maximum sequence length (optional) - disable_cuda_graph: Disable CUDA graphs (default: False) + - extra_args: Additional dynamo.sglang CLI flags (optional) num_workers: Number of SGLang worker processes single_gpu: If True, all workers share GPU 0 data_parallel_size: If set, enables this many data-parallel ranks per worker process. @@ -148,6 +152,7 @@ def __init__( max_total_tokens = sglang_args.get("max_total_tokens") context_length = sglang_args.get("context_length") disable_cuda_graph = sglang_args.get("disable_cuda_graph", False) + extra_args = sglang_args.get("extra_args", ()) # Resolved memory budget, for startup logs (mirrors the command flags). mem_budget = ( f"max_total_tokens={max_total_tokens}, mem_frac=0.9" @@ -216,6 +221,10 @@ def __init__( if disaggregation_mode is not None: command.extend(["--disaggregation-mode", disaggregation_mode]) command.extend(["--disaggregation-transfer-backend", "nixl"]) + # The KV bootstrap server binds to --host (default 127.0.0.1), + # but dynamo.sglang advertises the auto-detected local IP to + # decode workers, so bind all interfaces like launch/disagg.sh. + command.extend(["--host", "0.0.0.0"]) if data_parallel_size is not None: # Add DP configuration @@ -235,6 +244,7 @@ def __init__( kv_events_port = self._kv_event_ports[worker_idx * kv_event_rank_span] kv_events_config = f'{{"publisher":"zmq","topic":"kv-events","endpoint":"tcp://*:{kv_events_port}"}}' command.extend(["--kv-events-config", kv_events_config]) + command.extend(extra_args) nccl_port = self._nccl_ports[worker_idx] command.extend(["--nccl-port", str(nccl_port)]) @@ -417,7 +427,6 @@ def test_router_decisions_sglang_dp( ) -@pytest.mark.skip(reason="Nightly CI failure: https://linear.app/nvidia/issue/DYN-2603") @pytest.mark.e2e @pytest.mark.model(MODEL_NAME) @pytest.mark.gpu_2 @@ -451,6 +460,122 @@ def test_router_decisions_sglang_disagg( "gpu_start_index": 1, "disaggregation_mode": "decode", }, + test_kwargs={"enable_bootstrap": True}, + ) + + +# EAGLE3 speculative decoding: same pair as launch/agg_spec_decoding.sh, both +# ungated. An 8B base is needed because the test requires a real EAGLE3 draft +# that accepts tokens, and there is no smaller ungated base+draft pair. The two +# prefill workers share GPU0 (peak ~36 GiB), so the test is marked h100 and runs +# in the nightly sglang-h100-test lane. Like the other gpu_2 tests here it has no +# profiled_vram_gib, so it runs in that lane's sequential stage. +EAGLE_MODEL_NAME = "Qwen/Qwen3-8B" +EAGLE_DRAFT_MODEL_NAME = "Tengyunw/qwen3_8b_eagle3" +SGLANG_EAGLE_ARGS: Dict[str, Any] = { + **SGLANG_ARGS, + "model": EAGLE_MODEL_NAME, + "extra_args": ( + "--enable-metrics", + "--speculative-algorithm", + "EAGLE3", + "--speculative-draft-model-path", + EAGLE_DRAFT_MODEL_NAME, + "--speculative-num-steps", + "3", + "--speculative-eagle-topk", + "1", + "--speculative-num-draft-tokens", + "4", + ), +} + + +def _assert_decode_workers_speculated( + prefill_workers: SGLangProcess, decode_workers: SGLangProcess +) -> None: + """Assert every decode worker ran EAGLE verify steps and accepted draft tokens.""" + for port in decode_workers._system_ports: + response = requests.get(f"http://localhost:{port}/metrics", timeout=10) + response.raise_for_status() + content = response.text + SGLangSpecDecodeMetricsPayload( + body={}, + repeat_count=1, + expected_log=[], + expected_response=[], + min_num_requests=4, + ).validate(None, content) + + +@pytest.mark.e2e +@pytest.mark.model(EAGLE_MODEL_NAME) +@pytest.mark.model(EAGLE_DRAFT_MODEL_NAME) +@pytest.mark.h100 +@pytest.mark.gpu_2 +@pytest.mark.nightly +@pytest.mark.requested_sglang_kv_tokens(2048) +@pytest.mark.parametrize("request_plane", ["nats"], indirect=True) +# ~3x ~96s (2x H200, models pre-cached); peak ~36 GiB on the shared prefill GPU. +@pytest.mark.timeout(300) +def test_router_decisions_sglang_disagg_eagle( + request, + runtime_services_dynamic_ports, + predownload_models, + set_ucx_tls_no_mm, + request_plane, +): + """EAGLE3 on both prefill and decode workers, behind the KV router. + + On top of the disagg routing checks, this asserts that: + * the router predicts a KV hit on each shared-prefix request, i.e. the + prefill worker's EAGLE (bigram) KV event hashes match the router's + * the decode worker keeps speculating after the KV and EAGLE state + transfer, and accepts draft tokens + """ + run_disagg_router_decisions_test( + engine_process_cls=SGLangProcess, + engine_args_name="sglang_args", + engine_args=SGLANG_EAGLE_ARGS, + request=request, + request_plane=request_plane, + model_name=EAGLE_MODEL_NAME, + block_size=PAGE_SIZE, + num_prefill_workers=2, + num_decode_workers=1, + prefill_process_kwargs={ + "single_gpu": True, + "gpu_start_index": 0, + "disaggregation_mode": "prefill", + }, + decode_process_kwargs={ + "single_gpu": True, + "gpu_start_index": 1, + "disaggregation_mode": "decode", + }, + # Enough decode steps for a stable tokens-per-verify ratio. + test_payload={ + **build_test_payload(EAGLE_MODEL_NAME), + "max_tokens": 128, + "temperature": 0.0, + }, + # enable_bootstrap=True skips the kv_transfer_estimated_latency_ms > 0 + # check, which does not hold for SGLang. SGLang disagg always uses the + # bootstrap flow: + # 1. The router sends the request to a prefill worker. + # 2. The prefill worker replies at once with its bootstrap address, + # before prefill has finished. + # 3. The router dispatches the request to a decode worker, so + # prefill and decode now run in parallel. + # 4. The decode worker pulls the KV cache from the prefill worker and + # emits its first token. + # 5. The prefill worker's response stream ends. Only now does the + # router record prefill completion. + # The latency is (decode first token time - prefill completion time), + # clamped at 0. Steps 4 and 5 race, and when the first token arrives + # first the latency is exactly 0, so the > 0 check fails intermittently. + test_kwargs={"enable_bootstrap": True, "require_kv_hit": True}, + post_check=_assert_decode_workers_speculated, )