recipes/docker-images/mia-tp4/adapter/folded_result_fence.py

"""Detach folded DSpark verify results from the verify graph's persistent buffers. Default OFF.

DSV41_FOLDED_FENCE=1. On the folded path (all-greedy batches replayed from the verify CUDA graph),
accept_and_finalize returns views of DsparkVerifyEpilogue's persistent buffers (out_tokens_buf,
commit_lens_buf, ...). With overlap scheduling the scheduler copies them to host on its copy stream
while the next step already runs, and the next verify replay can overwrite those buffers before the
copy finishes (sglang#40919). Cloning them on the forward stream right after the verify, before
anything else is queued, gives the copy its own memory; later replays cannot touch it. Six tensors
of bs x 6 ints: negligible. Sampled batches take the eager path, which already allocates fresh
tensors, and are left alone.
"""
import os

ENABLED = os.environ.get("DSV41_FOLDED_FENCE", "0").strip() not in ("0", "", "off", "false")
FIELDS = ("correct_len", "bonus", "cap_trim_lens", "commit_lens", "new_seq_lens", "out_tokens")


def install(module):
    """sglang.srt.speculative.dspark_components.dspark_verify"""
    if not ENABLED:
        return
    cls = getattr(module, "TargetVerifyExecutor", None)
    outs_cls = getattr(module, "AcceptOuts", None)
    if cls is None or outs_cls is None or not hasattr(cls, "accept_and_finalize"):
        raise RuntimeError("DSV41_FOLDED_FENCE: TargetVerifyExecutor.accept_and_finalize / AcceptOuts gone; engine drifted")
    if getattr(cls, "_dsv41_folded_fence", False):
        return
    cls._dsv41_folded_fence = True
    orig = cls.accept_and_finalize

    def accept_and_finalize(self, *a, **kw):
        outs = orig(self, *a, **kw)
        if not kw.get("folded_accept", False):
            return outs
        return outs_cls(**{f: getattr(outs, f).clone() for f in FIELDS})

    cls.accept_and_finalize = accept_and_finalize
    print("[folded_fence] folded DSpark results cloned off the persistent verify buffers", flush=True)