minke.yu/sglang
1
0
Fork 0
Code Issues Pull Requests Actions 12 Packages Projects Releases Wiki Activity

[unified-memory] Enable prefill cuda-graph capture (#37418)

Browse Source
This commit is contained in:
Cheng Wan
2026-09-13 19:17:19 -07:00
committed by GitHub
parent 2ec4bbcbd4
commit 6410800af9
5 changed files with 154 additions and 60 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/registered/e2e/models/test_inkling_unified.py
+36 -10
View File
@@ -51,7 +51,7 @@ _MODEL_PATH = os.environ.get("INKLING_TEST_MODEL_PATH", "thinkingmachines/Inklin
_MODEL_REVISION = os.environ.get("INKLING_TEST_MODEL_REVISION", "test")
def _unified_args():
def _unified_args(*, attention_backend="triton", prefill_cuda_graph=False):
"""Server args for the tri-pool boot. Mirrors test_inkling.py's fixture
minus the multimodal/parser surface (KV-path focus), plus the unified
flags. The ratios still feed boot sizing until the byte configurator
@@ -59,17 +59,12 @@ def _unified_args():
args = [
"--trust-remote-code",
"--enable-unified-memory",
# Unified requires the Triton strided page-major read/write paths.
"--attention-backend",
"triton",
attention_backend,
"--page-size",
"128",
"--mamba-radix-cache-strategy",
"extra_buffer",
# Inkling defaults to a FULL prefill graph, which unified rejects at
# boot: the prefill graph runner bypasses the virtual->physical rebind.
"--cuda-graph-backend-prefill",
"disabled",
"--swa-full-tokens-ratio",
"0.1",
"--mamba-full-memory-ratio",
@@ -77,6 +72,11 @@ def _unified_args():
"--mem-fraction-static",
"0.5",
]
if not prefill_cuda_graph:
# Inkling declares a FULL prefill graph as a model default; the Triton
# cells cannot serve it (the cuda-graph metadata path has no EXTEND
# branch), so pin it off rather than lean on the auto-fallback.
args += ["--cuda-graph-backend-prefill", "disabled"]
if _MODEL_REVISION:
args += ["--revision", _MODEL_REVISION]
return args
@@ -91,8 +91,8 @@ def _static_args():
"128",
"--mamba-radix-cache-strategy",
"extra_buffer",
# Inkling defaults to a FULL prefill graph, which unified rejects at
# boot: the prefill graph runner bypasses the virtual->physical rebind.
# Match the unified cell's Triton pin, which cannot serve Inkling's
# default FULL prefill graph.
"--cuda-graph-backend-prefill",
"disabled",
"--swa-full-tokens-ratio",
@@ -128,6 +128,10 @@ def _greedy_generate(base_url, text, max_new_tokens=32, logprobs=False):
class TestInklingUnifiedTriPool(CustomTestCase):
@classmethod
def server_args(cls):
return _unified_args()
@classmethod
def setUpClass(cls):
cls.model = _MODEL_PATH
@@ -136,7 +140,7 @@ class TestInklingUnifiedTriPool(CustomTestCase):
cls.model,
cls.base_url,
timeout=DEFAULT_TIMEOUT_FOR_SERVER_LAUNCH,
other_args=_unified_args(),
other_args=cls.server_args(),
env={**os.environ, "SGLANG_ENABLE_UNIFIED_RADIX_TREE": "1"},
)
@@ -194,6 +198,28 @@ class TestInklingUnifiedTriPool(CustomTestCase):
self.assertGreater(len(data["text"].strip()), 0, data)
class TestInklingUnifiedFullPrefillGraph(TestInklingUnifiedTriPool):
"""The same tri-pool guards with Inkling's OWN default FULL prefill cuda
graph left on, over fa4 -- the backend family whose
`_init_full_cg_prefill_metadata` implements that path.
The pairing used to be a hard boot failure: unified disabled prefill
capture outright. With capture on, both the captured block table and the
SWA write loc have to come from the translator. Re-running the full->swa
map on `out_cache_loc` does not work here -- it is already FULL-side
kernel-facing by then, and indexes far past the swa v2p table (a
device-side "index out of bounds" assert).
`test_input_output_logprobs_match` is the sharp guard: a wrong-slot SWA
write moves logprobs at once, and
`test_long_decode_slides_past_swa_window` keeps compaction running
underneath a replaying graph.
"""
@classmethod
def server_args(cls):
return _unified_args(attention_backend="fa4", prefill_cuda_graph=True)
@unittest.skipUnless(
os.environ.get("INKLING_UNIFIED_PARITY") == "1",
"eval-host lane: set INKLING_UNIFIED_PARITY=1 (two sequential server boots)",
test/registered/unit/layers/attention/test_flashattention_graph_metadata.py
+11
View File
@@ -24,6 +24,17 @@ class TestFlashAttentionGraphMetadata(CustomTestCase):
backend.token_to_kv_pool = SimpleNamespace(
translate_loc_from_full_to_swa=lambda locations: locations
)
# The metadata builder reads `is_translating` to choose between the
# translated block table and the strided one this test covers, so the
# source has to be real; the stub pools disable translation, which is
# the static-pool view the assertions below are written against.
backend.kv_index_translator = KVIndexTranslator(
req_to_token=backend.req_to_token,
token_to_kv_pool_allocator=SimpleNamespace(),
token_to_kv_pool=SimpleNamespace(),
page_size=backend.page_size,
device="cpu",
)
forward_batch = SimpleNamespace(
batch_size=1,
seq_lens=torch.zeros(1, dtype=torch.int64),
test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py
+49 -30
View File
@@ -11,26 +11,31 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""`--enable-unified-memory` disables PREFILL cuda-graph capture.
"""`--enable-unified-memory` and PREFILL cuda-graph capture.
BUG REGRESSION. Only decode capture is wired: the prefill graph runner builds
its ForwardBatch directly, so it never runs the unified pool's write-loc
rebind (rebind_write_loc) and the captured batch holds VIRTUAL ids -- the
captured store would silently write wrong slots.
Capture is wired: the captured batch reads `out_cache_loc` out of the registry
slot, refilled before each replay from the already-rebound kernel-facing loc,
and the read tables are refilled out-of-graph from the live v2p. So BREAKABLE
(the CUDA default) and TC_PIECEWISE must be left alone -- an earlier gate
disabled every prefill backend outright, which cost every unified run its
prefill graph.
The old gate only rejected `TC_PIECEWISE`, but the generic prefill default is
`BREAKABLE` -- so the DEFAULT unified invocation was broken; it only ever
worked when `--cuda-graph-backend-prefill=disabled` happened to be passed.
The FULL backend is the exception, and for a reason that is not about unified
memory: its metadata path (`_init_full_cg_prefill_metadata`) is implemented
only by the fa3/fa4 family. Anything else lands in the decode-shaped
`_apply_cuda_graph_metadata`, which has no EXTEND branch.
Pinned: the default is auto-disabled with a warning (unified boots out of the
box), an EXPLICIT prefill backend still raises (never silently override a
user's stated intent), and decode capture is untouched either way.
Pinned here: FULL survives on fa3/fa4, FULL is disabled with a warning on any
other backend (Inkling declares FULL as a MODEL default, so refusing to boot
would fail on a flag the user never typed), and decode capture is never
touched.
python -m pytest test/registered/unit/server_args/test_unified_prefill_cuda_graph_gate.py -v
"""
import unittest
from types import SimpleNamespace
from unittest.mock import patch
import msgspec
@@ -42,7 +47,7 @@ from sglang.test.ci.ci_register import register_cpu_ci
register_cpu_ci(est_time=8, suite="base-a-test-cpu")
def _run_handler(*, prefill_backend, explicit):
def _run_handler(*, prefill_backend, attention_backends):
"""Run just `handle_unified_memory_pool` over a minimal stand-in."""
sa = ServerArgs(model_path="dummy")
cg = SimpleNamespace(
@@ -59,33 +64,47 @@ def _run_handler(*, prefill_backend, explicit):
"enable_two_batch_overlap": False,
"dcp_size": 1,
"cuda_graph_config": cg,
"cuda_graph_backend_prefill": prefill_backend if explicit else None,
"cuda_graph_backend_prefill": prefill_backend,
}.items():
msgspec.Struct.__setattr__(sa, name, value)
handle_unified_memory_pool(sa)
with patch(
"sglang.srt.arg_groups.kv_cache_hook.attention_backends_of",
return_value=attention_backends,
):
handle_unified_memory_pool(sa)
return cg
class TestUnifiedPrefillCudaGraphGate(unittest.TestCase):
def test_default_prefill_capture_is_auto_disabled(self):
"""The generic default (BREAKABLE) must be turned off, not crash the
server 30 seconds later inside graph capture."""
for backend in (Backend.BREAKABLE, Backend.FULL, Backend.TC_PIECEWISE):
cg = _run_handler(prefill_backend=backend, explicit=False)
self.assertEqual(cg.prefill.backend, Backend.DISABLED)
# Decode capture is the wired path and must survive untouched.
self.assertEqual(cg.decode.backend, Backend.FULL)
def test_non_full_prefill_backends_are_left_enabled(self):
"""BUG REGRESSION. Unified used to disable prefill capture outright, so
the default BREAKABLE graph silently never ran."""
for backend in (Backend.BREAKABLE, Backend.TC_PIECEWISE):
for attn in (("fa4", "fa4"), ("triton", "triton")):
with self.subTest(prefill=backend, attn=attn):
cg = _run_handler(prefill_backend=backend, attention_backends=attn)
self.assertEqual(cg.prefill.backend, backend)
self.assertEqual(cg.decode.backend, Backend.FULL)
def test_explicit_prefill_backend_is_refused(self):
"""A user who explicitly asked for prefill graphs gets a clear error,
not a silent override of their stated intent."""
for backend in (Backend.BREAKABLE, Backend.FULL, Backend.TC_PIECEWISE):
with self.assertRaises(ValueError) as ctx:
_run_handler(prefill_backend=backend, explicit=True)
self.assertIn("prefill capture is not wired", str(ctx.exception))
def test_full_prefill_survives_on_the_fa_family(self):
for attn in (("fa3", "fa3"), ("fa4", "fa4")):
with self.subTest(attn=attn):
cg = _run_handler(prefill_backend=Backend.FULL, attention_backends=attn)
self.assertEqual(cg.prefill.backend, Backend.FULL)
def test_full_prefill_is_disabled_on_other_backends(self):
"""Warn and fall back rather than raise: Inkling declares FULL as a
model default, indistinguishable at this point from a user flag."""
for attn in (("triton", "triton"), ("flashinfer", "flashinfer")):
with self.subTest(attn=attn):
cg = _run_handler(prefill_backend=Backend.FULL, attention_backends=attn)
self.assertEqual(cg.prefill.backend, Backend.DISABLED)
self.assertEqual(cg.decode.backend, Backend.FULL)
def test_already_disabled_is_a_no_op(self):
cg = _run_handler(prefill_backend=Backend.DISABLED, explicit=True)
cg = _run_handler(
prefill_backend=Backend.DISABLED, attention_backends=("triton", "triton")
)
self.assertEqual(cg.prefill.backend, Backend.DISABLED)
self.assertEqual(cg.decode.backend, Backend.FULL)