diff --git a/python/freetoken/engine/graph.py b/python/freetoken/engine/graph.py index a995316ddc..0d79bad93f 100644 --- a/python/freetoken/engine/graph.py +++ b/python/freetoken/engine/graph.py @@ -10,6 +10,7 @@ import torch from freetoken.core import Batch, Req, get_global_ctx from freetoken.distributed import get_tp_info +from freetoken.env import ENV from freetoken.utils import init_logger, mem_GB from freetoken.utils.progress import emit_progress from tqdm import tqdm @@ -120,7 +121,7 @@ def _upload_graph(graph: torch.cuda.CUDAGraph, stream: torch.cuda.Stream) -> Non Otherwise the first replay pushes every node behind the graph's own stream WAIT, and a large graph then blocks cudaGraphLaunch until the WAIT is released. PLE wait-sync releases it only after the launch returns, so that deadlocks. """ - upload = _cu_graph_upload() + upload = _cu_graph_upload() if ENV.GRAPH_UPLOAD else None if upload is None: return rc = upload(graph.raw_cuda_graph_exec(), stream.cuda_stream) diff --git a/python/freetoken/env.py b/python/freetoken/env.py index 9a06b1c029..03e9388024 100644 --- a/python/freetoken/env.py +++ b/python/freetoken/env.py @@ -77,6 +77,9 @@ class EnvClassSingleton: MAMBA_SSM_DTYPE = EnvStr("float32") # cap the longest prefill warmup forward; 0 = the engine's own chunk cap WARMUP_MAX_LEN = EnvInt(0) + # cuGraphUpload each decode graph at capture (#697); off keeps the first-launch upload, which + # replays ~0.2 ms/step faster on Qwen3.8-Flash / RTX 5080 but can block the first launch on a stream WAIT + GRAPH_UPLOAD = EnvBool(True) # GPU-decode MoE: run a layer's expert miss copy on a side stream under the shared expert MOE_COPY_OVERLAP = EnvBool(False) # GPU-decode MoE cross-layer prefetch: off | measure (predict and count, no copies) | on (predict and copy) diff --git a/tests/engine/test_graph_upload.py b/tests/engine/test_graph_upload.py new file mode 100644 index 0000000000..de61152162 --- /dev/null +++ b/tests/engine/test_graph_upload.py @@ -0,0 +1,28 @@ +import pytest + +import freetoken.engine.graph as graph_mod +from freetoken.env import ENV + + +class _Graph: + def raw_cuda_graph_exec(self): + return 1234 + + +class _Stream: + cuda_stream = 5678 + synced = False + + def synchronize(self): + self.synced = True + + +@pytest.mark.parametrize("enabled", [True, False]) +def test_graph_upload_follows_env(monkeypatch, enabled): + calls = [] + monkeypatch.setattr(graph_mod, "_cu_graph_upload", lambda: lambda exec_, stream: calls.append((exec_, stream)) or 0) + monkeypatch.setattr(ENV.GRAPH_UPLOAD, "value", enabled) + stream = _Stream() + graph_mod._upload_graph(_Graph(), stream) + assert calls == ([(1234, 5678)] if enabled else []) + assert stream.synced is enabled