[Model] Support Ling-3.0-flash (BailingMoeV3) (#33561)

Signed-off-by: JustinTong <justintong0323@gmail.com>
Signed-off-by: Xinyuan Tong <xinyuantong.cs@gmail.com>
Co-authored-by: luoyuan.luo <luoyuan.luo@antgroup.com>
Co-authored-by: 得泽 <zhangkaihong.zkh@antgroup.com>
Co-authored-by: 翎悦 <vito.yy@antgroup.com>
Co-authored-by: 羽癫 <yudian.zy@antgroup.com>
Co-authored-by: tiwei.btw <tiwei.btw@antgroup.com>
Co-authored-by: Liangsheng Yin <hnyls2002@gmail.com>
Co-authored-by: 文赋 <zibin.zb@antgroup.com>
Co-authored-by: JustinTong <justintong0323@gmail.com>

python/sglang/srt/model_executor/pool_configurator.py

@@ -213,7 +213,11 @@ class DefaultPoolConfigurator(MemoryPoolConfigurator):
                         self._cell_size * (1 + draft_num_layers / int(num_layers))
                     )
 
-        # DFLASH/DSPARK: scale cell_size to account for draft model KV cache
+        # DFLASH/DSPARK: reserve the draft runner's *actual* per-token KV cost.
+        # The draft allocates its own KV pool at the target's
+        # max_total_num_tokens, whose per-token footprint can differ from the
+        # target's (e.g. an MLA-latent target paired with a full per-head K/V
+        # draft), so size from the draft config rather than the layer ratio.
         if kvc.spec_algorithm.is_dflash_family() and not kvc.is_draft_worker:
             from sglang.srt.speculative.dflash_utils import (
                 scale_kv_cell_size_per_token_for_dflash,

python/sglang/srt/model_executor/runner/decode_cuda_graph_runner.py

@@ -78,6 +78,7 @@ from sglang.srt.model_executor.runner.base_cuda_graph_runner import (
 from sglang.srt.model_executor.runner.flashinfer_autotune import (
     maybe_flashinfer_autotune_speculative_draft,
 )
+from sglang.srt.model_executor.runner.metadata_glue_graph import MetadataGlueGraph
 from sglang.srt.model_executor.runner.shape_key import ShapeKey
 from sglang.srt.model_executor.runner_backend.breakable_cuda_graph_backend import (
     BreakableCudaGraphBackend,
@@ -457,6 +458,25 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             source=self.buffers,
         )
 
+        # Captures the per-replay attention-metadata prep into a small CUDA
+        # graph; see metadata_glue_graph.py for the correctness contract.
+        # Force-off for DFlash-family spec: verify installs host-fed fast
+        # plans (sync-free begin_forward that recomputes plan inputs on the
+        # host every replay), and capturing one freezes the capture-time
+        # plan — drafts go stale and accept length collapses to ~1.
+        enable_metadata_glue = envs.SGLANG_ENABLE_METADATA_GLUE_GRAPH.get()
+        if enable_metadata_glue and model_runner.spec_algorithm.is_dflash_family():
+            logger.warning(
+                "SGLANG_ENABLE_METADATA_GLUE_GRAPH is incompatible with "
+                "DFlash-family speculative decoding (host-fed fast verify "
+                "plans must re-run on the host every replay); disabling the "
+                "metadata glue graph."
+            )
+            enable_metadata_glue = False
+        self._metadata_glue = (
+            MetadataGlueGraph(self.device) if enable_metadata_glue else None
+        )
+
         # --- backend ---------------------------------------------------
         self.backend = resolve_decode_backend(self)
 
@@ -1367,7 +1387,34 @@ class DecodeCudaGraphRunner(BaseCudaGraphRunner):
             capture_forward_mode=self.capture_forward_mode,
             is_encoder_decoder=self.is_encoder_decoder,
         )
-        attn_backend.init_forward_metadata_out_graph(fb_view)
+        # Glue-graph fast path: pointer-stable prep (static buffers + pool
+        # tensors only) is captured per key; guards keep every python-visible
+        # branch inside the backends constant for that key.
+        if (
+            self._metadata_glue is not None
+            and not self._metadata_glue.disabled
+            and raw_bs == bs
+            and not self.enable_two_batch_overlap
+            and not self.enable_pdmux
+            and self.model_runner.lora_manager is None
+        ):
+            # actual_forward_mode belongs in the key even though the captured
+            # graph always targets capture_forward_mode: DSV4's replay prep
+            # substitutes seq_lens / seq_lens_cpu / seq_lens_sum /
+            # req_pool_indices / out_cache_loc when the runtime mode is IDLE,
+            # so IDLE and active DECODE are different python branches and must
+            # not share a captured graph.
+            self._metadata_glue.run(
+                attn_backend,
+                fb_view,
+                (
+                    bs,
+                    str(self.capture_forward_mode),
+                    str(fb_view.actual_forward_mode),
+                ),
+            )
+        else:
+            attn_backend.init_forward_metadata_out_graph(fb_view)
 
         self.raw_bs = raw_bs
         self.raw_num_token = raw_num_token

python/sglang/srt/model_executor/runner/metadata_glue_graph.py

@@ -0,0 +1,106 @@
+"""Glue-graph capture of the per-replay attention-metadata prep.
+
+``decode_cuda_graph_runner.load_batch`` runs
+``attn_backend.init_forward_metadata_out_graph(fb_view)`` eagerly on every
+replay. At bs=1 spec decode this is an "op soup": dozens of tiny tensor ops
+whose HOST dispatch cost dominates the inter-phase seam, while every device
+input/output lives at a stable address — the replay fb view hands backends the
+runner's static buffers, and pool tensors are persistent. Capturing the op
+sequence once per replay key collapses the per-step host cost to a single
+graph launch.
+
+Correctness contract:
+
+- The caller only routes here when the replay is padding-free
+  (raw_bs == padded bs) and TBO / pdmux / LoRA are off, so every
+  Python-visible branch inside the backends is constant per key.
+- Python side effects (each backend's ``forward_metadata`` object) are
+  snapshotted at capture time and re-installed on every replay; the graph
+  replays only the device ops that refresh the tensors those objects point to.
+- ``NUM_WARMUP`` eager runs precede capture so triton JIT compile / autotune
+  happen outside capture.
+- Any capture failure (e.g. a backend syncing or reading host values inside
+  its prep) permanently disables the glue graph and falls back to eager.
+- Backends whose prep computes values on the HOST each replay (e.g. the
+  DFlash-family host-fed fast verify plans) must never be glued: capture
+  records only device ops, so the host-written plan inputs would replay
+  frozen at their capture-time values. Note the failure is SILENT — capture
+  succeeds, outputs stay correct, only accept length collapses. Callers must
+  gate such configurations off before routing here
+  (``decode_cuda_graph_runner`` force-disables the glue for DFlash-family
+  spec).
+"""
+
+from __future__ import annotations
+
+import logging
+from typing import Any, Dict, List
+
+import torch
+
+logger = logging.getLogger(__name__)
+
+
+class MetadataGlueGraph:
+    NUM_WARMUP = 2
+
+    def __init__(self, device):
+        self.device = device
+        self.disabled = False
+        self._states: Dict[Any, dict] = {}
+        self._capture_stream = None
+
+    def reset(self):
+        """Drop captured graphs (call when the runner recaptures its graphs —
+        static buffers and backend state may have been rebuilt)."""
+        self._states.clear()
+
+    @staticmethod
+    def _leaves(attn_backend) -> List[Any]:
+        backends = [attn_backend]
+        if attn_backend.attn_backend_list is not None:
+            backends.extend(attn_backend.attn_backend_list)
+        return backends
+
+    def run(self, attn_backend, fb_view, key) -> None:
+        """Run ``init_forward_metadata_out_graph`` for this replay, through the
+        captured glue graph once it is ready."""
+        st = self._states.get(key)
+        if st is None:
+            st = {"warmups": 0, "graph": None, "meta": None}
+            self._states[key] = st
+
+        if st["graph"] is not None:
+            for backend, metadata in st["meta"]:
+                backend.forward_metadata = metadata
+            st["graph"].replay()
+            return
+
+        if st["warmups"] < self.NUM_WARMUP:
+            st["warmups"] += 1
+            attn_backend.init_forward_metadata_out_graph(fb_view)
+            return
+
+        if self._capture_stream is None:
+            self._capture_stream = torch.cuda.Stream()
+        graph = torch.cuda.CUDAGraph()
+        try:
+            with torch.cuda.graph(graph, stream=self._capture_stream):
+                attn_backend.init_forward_metadata_out_graph(fb_view)
+        except Exception:
+            logger.warning(
+                "Metadata glue-graph capture failed for key %s; falling back "
+                "to eager metadata prep permanently.",
+                key,
+                exc_info=True,
+            )
+            self.disabled = True
+            # Ops under a failed capture were recorded, not executed — run
+            # this step's prep for real.
+            attn_backend.init_forward_metadata_out_graph(fb_view)
+            return
+
+        st["meta"] = [(b, b.forward_metadata) for b in self._leaves(attn_backend)]
+        st["graph"] = graph
+        # Capture records without executing; replay once to do this step's prep.
+        graph.replay()