perf(engine): FREETOKEN_GRAPH_UPLOAD switch; #697's capture-time upload costs 1-2% decode

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)

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)

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