Skip to content

[HiCache][PD] fix: keep decode load-back polling out of rank-local collectives - #41393

Merged
ispobock merged 5 commits into
sgl-project:mainfrom
alphabetc1:fix-pd-hicache-loadback-collective
Oct 1, 2026
Merged

ispobock merged 5 commits into
sgl-project:mainfrom
alphabetc1:fix-pd-hicache-loadback-collective

Conversation

@alphabetc1

@alphabetc1 alphabetc1 commented Sep 27, 2026 •

Copy link
Copy Markdown
Collaborator

Motivation

PD decode with --disaggregation-decode-enable-radix-cache --enable-hierarchical-cache and TP > 1 crashes on the first L2 load-backs. _process_hicache_local_restores enters TP collectives based on each rank's own CUDA event:

  • Phase A: is_load_back_event_done ran loading_check() (a TP all_reduce) only on ranks whose event had fired.
  • Phase B: new load-backs were gated on the producer_index + 1 event, so the same all_reduce, plus init_load_back / start_loading, ran on a subset of ranks.

On Qwen3-8B, prefill TP1 + decode TP2, the decode scheduler dies with:

gloo::EnforceNotMet ... op.preamble.length <= op.nbytes. 8 vs 1. Received data size doesn't match expected size. Is there a distributed collective mismatch in your code?

With SGLANG_ENABLE_RANK_CONSENSUS_CHECKER=1 and probes on this path's collective entries (see Reproduction), the first divergence is:

[TP0] Found rank divergence ... events = ['loading_check(all_reduce) from _process_hicache_local_restores:285']
[TP1] Found rank divergence ... events = ['loading_check(all_reduce) from _process_hicache_local_restores:295']

TP0 entered from Phase A and TP1 from Phase B. The two same-shape all_reduces pair silently; the gloo crash is the next, differently sized collective. #29483 covered only Phase A, on the removed hiradix_cache.py.

Modifications

  • is_load_back_event_done only polls the local event. Load acks are already reaped in lockstep by check_hicache_events, and a rank-local READY is harmless because the gated KV poll is MIN-reduced.
  • Phase B gates on has_free_load_back_slot(): len(ack_load_queue) < num_counters. The queue grows in start_loading and shrinks by a MIN-reduced count, so all ranks agree, with no new collective. A reaped ack was synchronized, so the reused slot's events are done. A slot is reused at most one step later.

Reproduction

repro.py, sitecustomize.py, commands
# repro.py
import concurrent.futures as cf
import os
import random

import requests

from sglang.test.server_fixtures.disaggregation_fixture import PDDisaggregationServerBase


class Run(PDDisaggregationServerBase):
    prefill_tp_size = 1
    decode_tp_size = 2
    decode_base_gpu_id = 1
    extra_prefill_args = ["--mem-fraction-static", "0.8"]
    extra_decode_args = [
        "--disaggregation-decode-enable-radix-cache",
        "--enable-hierarchical-cache",
        "--hicache-ratio", "4",
        "--hicache-write-policy", "write_through",
        "--max-total-tokens", "24000",
        "--enable-metrics",
    ]
    extra_decode_env = {
        k: v for k, v in os.environ.items()
        if k.startswith(("REPRO_", "SGLANG_ENABLE_RANK"))
    }


def gen(ids):
    r = requests.post(
        f"{Run.base_url}/generate",
        json={"input_ids": ids, "sampling_params": {"max_new_tokens": 16, "temperature": 0, "ignore_eos": True}},
        timeout=180,
    )
    r.raise_for_status()
    return r.json()["meta_info"]["cached_tokens"]


Run.setUpClass()
Run.model = "Qwen/Qwen3-8B"
Run.launch_all()
try:
    rng = random.Random(0)
    prefixes = [[rng.randint(1000, 30000) for _ in range(3000)] for _ in range(24)]
    with cf.ThreadPoolExecutor(32) as ex:
        for rnd in range(8):
            jobs = [ex.submit(gen, p + [rng.randint(1000, 30000) for _ in range(32)]) for p in prefixes]
            print(f"round={rnd} cached={sum(j.result() for j in jobs)}", flush=True)
finally:
    Run.tearDownClass()
# sitecustomize.py
# REPRO_H2D_SLEEP_MS: delay rank 1's H2D stream so the ranks see the same load-back event at different times.
# REPRO_PROBES=1: record every collective entry on this path with the rank consensus checker; the sleep lets
#   the checker's worker compare events before the mismatched collective aborts the process.
import os
import sys
import time

MS = float(os.environ.get("REPRO_H2D_SLEEP_MS", "0"))
PROBES = os.environ.get("REPRO_PROBES") == "1"
if MS > 0 or PROBES:
    import torch
    from sglang.srt.disaggregation import decode
    from sglang.srt.managers import cache_controller as cc
    from sglang.srt.mem_cache import unified_radix_cache as urc
    from sglang.srt.utils import rank_consensus_checker as rdc

    orig_start_loading = cc.HiCacheController.start_loading

    def start_loading(self):
        if MS > 0 and self.load_queue and torch.distributed.get_rank() == 1:
            with torch.cuda.stream(self.l2_transfer_engine.host_to_device_stream):
                torch.cuda._sleep(int(MS * 1.98e6))
        if PROBES:
            rdc.assert_same("start_loading qlen=%d", len(self.load_queue))
        return orig_start_loading(self)

    cc.HiCacheController.start_loading = start_loading

    if PROBES:
        orig_loading_check = urc.UnifiedRadixCache.loading_check

        def loading_check(self, finish_count=None):
            if finish_count is None:
                f = sys._getframe(2)
                rdc.assert_same("loading_check(all_reduce) from %s:%d", f.f_code.co_name, f.f_lineno)
                time.sleep(2)
            return orig_loading_check(self, finish_count)

        urc.UnifiedRadixCache.loading_check = loading_check

        orig_poll = decode.poll_and_all_reduce

        def poll_and_all_reduce(pollers, *args, **kwargs):
            rdc.assert_same("poll_and_all_reduce n=%d", len(pollers))
            return orig_poll(pollers, *args, **kwargs)

        decode.poll_and_all_reduce = poll_and_all_reduce
python repro.py   # main on H200: gloo crash in round 1-2
# force divergent events, and let the checker name the divergence:
SGLANG_ENABLE_RANK_CONSENSUS_CHECKER=1 REPRO_H2D_SLEEP_MS=30 REPRO_PROBES=1 \
  PYTHONPATH=<dir of sitecustomize.py>:$PYTHONPATH python repro.py

Without REPRO_PROBES the checker does not catch it: loading_check and poll_and_all_reduce carry no @rank_consensus, so the gloo abort comes first.

Accuracy Tests

4xH200, 8 rounds per run, rounds 1-7 served entirely from L2 (sglang:load_back_tokens_total = 72,000 per rank per round):

main this PR
natural, checker on gloo crash healthy
30 ms H2D sleep on rank 1, checker on gloo crash healthy
+ probes Found rank divergence above healthy, checker silent

A temporary 2-rank Gloo test with rank-divergent events failed on main and when either Phase A or Phase B was reverted alone, and passed with this PR (not kept in the diff).

Checklist

  • Format your code according to the Format code with pre-commit.
  • No test files are changed in the final diff; validation is recorded above.

CI States

Latest PR Test (Base): ✅ Run #36712962654
Latest PR Test (Extra): 🚫 Run #36712962296
Latest PR Test (AMD ROCm 10): ❌ Run #36712962617

@alphabetc1 alphabetc1 added run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) labels Sep 29, 2026
@ispobock
ispobock merged commit b62b21e into sgl-project:main Oct 1, 2026
350 of 406 checks passed
@alphabetc1
alphabetc1 deleted the fix-pd-hicache-loadback-collective branch October 1, 2026 04:15
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

bypass-fail-fast CI: a failing job no longer aborts its siblings (lint still gates) hicache Hierarchical Caching for SGLang run-ci CI: run the baseline test suite on this PR run-ci-extra CI: also run the extra suite (requires run-ci) unified-radix-cache

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants