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

python/sglang/srt/arg_groups/kv_cache_hook.py

@@ -262,25 +262,34 @@ def handle_unified_memory_pool(server_args: Any) -> None:
     )
     if cfg.dcp_size > 1:
         _validate_unified_memory_dcp(server_args)
-    # Only monolithic decode cuda-graph capture is wired; piecewise prefill
-    # capture is not. Guard when the user opts into it.
+    # Prefill cuda-graph capture IS wired for the unified pool: the captured
+    # batch reads `out_cache_loc` out of the registry slot, which
+    # `populate_from_forward_batch` refills from the already-rebound (kernel-
+    # facing) loc before every replay, and the read tables are refilled
+    # out-of-graph from the live v2p.
+    #
+    # The FULL backend is the one exception, and not for a unified reason: its
+    # metadata path (`_init_full_cg_prefill_metadata`) exists only on the
+    # fa3/fa4 family. Any other backend lands in the decode-shaped
+    # `_apply_cuda_graph_metadata`, which has no EXTEND branch at all. Inkling
+    # declares FULL as a MODEL default, indistinguishable here from a flag the
+    # user typed, so warn and fall back rather than refuse to boot.
     _cg_cfg = cfg.cuda_graph_config
-    if _cg_cfg is not None and _cg_cfg.prefill.backend != Backend.DISABLED:
-        if cfg.cuda_graph_backend_prefill is not None:
-            raise ValueError(
-                "--enable-unified-memory supports decode cuda-graph "
-                "capture only; prefill capture is not wired (the prefill "
-                "graph runner bypasses the unified virtual->physical loc "
-                "rebind). Got --cuda-graph-backend-prefill="
-                f"{cfg.cuda_graph_backend_prefill!r}; pass "
-                "--cuda-graph-backend-prefill=disabled."
+    if _cg_cfg is not None and _cg_cfg.prefill.backend == Backend.FULL:
+        full_cg_backends = {"fa3", "fa4"}
+        backends = set(attention_backends_of(resolved_view(server_args)))
+        backends.discard(None)
+        if not backends <= full_cg_backends:
+            _cg_cfg.prefill.backend = Backend.DISABLED
+            logger.warning(
+                "--enable-unified-memory: disabling the FULL prefill "
+                "cuda-graph backend. It builds its block table in "
+                "_init_full_cg_prefill_metadata, which only %s implement; the "
+                "resolved attention backends are %s. Decode capture and the "
+                "other prefill backends are unaffected.",
+                sorted(full_cg_backends),
+                sorted(backends),
             )
-        _cg_cfg.prefill.backend = Backend.DISABLED
-        logger.warning(
-            "--enable-unified-memory: disabling prefill cuda-graph "
-            "capture (not wired for the unified pool's loc rebind); "
-            "decode capture is unaffected."
-        )
 
 
 def _validate_unified_memory_dcp(server_args: Any) -> None:

python/sglang/srt/layers/attention/flashattention_backend.py

@@ -26,6 +26,7 @@ from sglang.srt.layers.attention.verify_mask import VerifyMask, maybe_create_ver
 from sglang.srt.layers.cp.base import CPAttentionBackendKind, get_cp_strategy
 from sglang.srt.layers.cp.utils import is_cp_active
 from sglang.srt.layers.radix_attention import AttentionType
+from sglang.srt.mem_cache.kv_index_translator import KVReadTables
 from sglang.srt.mem_cache.memory_pool import KVWriteLoc
 from sglang.srt.mem_cache.swa_memory_pool import SWAKVPool
 from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode
@@ -673,8 +674,26 @@ class FlashAttentionBackend(AttentionBackend):
         m.cu_seqlens_q[1:].copy_(
             torch.cumsum(forward_batch.extend_seq_lens[:bs], dim=0)
         )
+        translating = self.kv_index_translator.is_translating
         max_seq_len_k = int(forward_batch.seq_lens_cpu[:bs].max().item())
-        if max_seq_len_k > 0:
+        if translating:
+            # Unified pool: the block table is a TRANSLATED page table, built
+            # straight into these capture-stable buffers from the LIVE v2p, so
+            # a page relocated by compaction since capture is picked up. Same
+            # substitution the eager extend branch makes in its `_unified_read`
+            # fixup; `build_index_table` emits page-granular kernel-facing ids
+            # directly, so there is no `// page_size` to undo.
+            self.kv_index_translator.build_index_table(
+                req_pool_indices=forward_batch.req_pool_indices[:bs],
+                seq_lens=forward_batch.seq_lens[:bs],
+                into=KVReadTables(
+                    full=m.page_table,
+                    sliding_window=(
+                        m.swa_page_table if self.use_sliding_window_kv_pool else None
+                    ),
+                ),
+            )
+        elif max_seq_len_k > 0:
             # Build the block table like the eager extend branch: take every
             # page_size-th token slot from req_to_token and divide by page_size.
             # Identity for page_size == 1 (strided is 0..max_seq_len_k-1, //1).
@@ -700,11 +719,21 @@ class FlashAttentionBackend(AttentionBackend):
                 self.full_cg_prefill_swa_out_cache_loc.shape[0],
                 "full-CG prefill SWA write-location buffer",
             )
-            self.full_cg_prefill_swa_out_cache_loc[:num_out].copy_(
-                self.token_to_kv_pool.translate_loc_from_full_to_swa(
+            # Under the unified pool `out_cache_loc` was rebound to FULL-side
+            # KERNEL-FACING ids at ForwardBatch construction, so the full->swa
+            # map cannot be re-run on it -- those values index far past the swa
+            # v2p table (a device-side "index out of bounds" assert). Phase 2 of
+            # the write contract derives the swa loc from them instead.
+            swa_write_loc = (
+                self.kv_index_translator.sliding_window_write_loc_for(
+                    forward_batch.out_cache_loc
+                )
+                if translating
+                else self.token_to_kv_pool.translate_loc_from_full_to_swa(
                     forward_batch.out_cache_loc
                 )
             )
+            self.full_cg_prefill_swa_out_cache_loc[:num_out].copy_(swa_write_loc)
             # Captured kernels read the full bucket. Route its inactive tail to
             # SWA's zero dummy slot to prevent stale writes into live slots.
             self.full_cg_prefill_swa_out_cache_loc[num_out:].zero_()