Stabilize spec-v2 draft graph metadata

Spec-v2 draft extend can receive token ids from producers whose dtype is not already int64, while DP collective paths require a stable integer dtype across ranks. EAGLE draft CUDA graph replay also pads raw batches to a captured batch size, so the metadata/replay path must see seq_lens_sum consistent with the padded seq_lens and then restore the caller-visible raw value.

Constraint: Keep this as a narrow correctness port from upstream rather than pulling the larger spec-v2 refactor chain.

Rejected: Cherry-pick broader attention-backend and decode-result refactors | current branch lacks the same upstream forward-context scaffolding and would require a separate port.

Confidence: high

Scope-risk: narrow

Directive: Do not remove the seq_lens_sum restore without rechecking padded EAGLE draft CUDA graph metadata construction.

Tested: python -m pytest test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py -q

Tested: remote g0034/cjy-glm5-new PYTHONPATH=python python3 -m pytest test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py -q

Not-tested: full multi-node GLM5 spec-v2 decode startup smoke

Co-authored-by: OmX <omx@oh-my-codex.dev>

python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py

@@ -411,12 +411,19 @@ class EAGLEDraftCudaGraphRunner:
             buffers.seq_lens_cpu[:raw_bs].copy_(forward_batch.seq_lens_cpu)
             forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:bs]
 
+        # Save the raw seq_lens_sum and keep it consistent with padded seq_lens
+        # while replay metadata and graph kernels observe the padded fake rows.
+        raw_seq_lens_sum = forward_batch.seq_lens_sum
+        if bs != raw_bs and raw_seq_lens_sum is not None:
+            forward_batch.seq_lens_sum = raw_seq_lens_sum + (
+                bs - raw_bs
+            ) * self.seq_len_fill_value
+
         self.model_runner.draft_attn_backend.init_forward_metadata_replay_cuda_graph(
             forward_batch, bs
         )
         self.raw_bs = raw_bs
         self.bs = bs
-        # TODO: The forward_batch.seq_len_sum might need to be updated to reflect the padding in the cuda graph
 
         # Replay
         self._replay(forward_batch)
@@ -430,5 +437,6 @@ class EAGLEDraftCudaGraphRunner:
             forward_batch.req_pool_indices = buffers.req_pool_indices[:raw_bs]
             if forward_batch.seq_lens_cpu is not None:
                 forward_batch.seq_lens_cpu = buffers.seq_lens_cpu[:raw_bs]
+            forward_batch.seq_lens_sum = raw_seq_lens_sum
 
         return out

python/sglang/srt/speculative/eagle_info_v2.py

@@ -192,7 +192,10 @@ class EagleDraftInputV2Mixin:
         extend_num_tokens = len(batch.seq_lens) * num_draft_tokens
 
         batch.spec_info = self
-        batch.input_ids = predict
+        # Normalize draft token ids before ForwardBatch construction; DP
+        # collectives require input_ids to have a consistent integer dtype
+        # across ranks.
+        batch.input_ids = predict.to(torch.int64)
         batch.extend_seq_lens = [num_draft_tokens for _ in range(len(batch.seq_lens))]
         batch.extend_prefix_lens = seq_lens_cpu_.tolist()
         batch.extend_num_tokens = extend_num_tokens

test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py

@@ -56,6 +56,88 @@ def test_eagle_v2_draft_extend_prepare_does_not_advance_source_batch_lengths():
     assert ("forward_batch", "seq_lens_sum") in assigned
 
 
+def test_eagle_v2_draft_extend_input_ids_are_normalized_to_int64():
+    """Draft input IDs must have a stable integer dtype before ForwardBatch init.
+
+    DeepSeek/GLM DP collectives assume `input_ids` dtype is consistent across
+    ranks.  Leaving this as the raw `predict` tensor lets an int32 producer leak
+    into the draft extend path.
+    """
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_info_v2.py")
+    func = _find_function(tree, "prepare_for_extend_to_fill_draft_kvcache")
+
+    for node in ast.walk(func):
+        if not isinstance(node, ast.Assign):
+            continue
+        for target in node.targets:
+            if (
+                isinstance(target, ast.Attribute)
+                and isinstance(target.value, ast.Name)
+                and target.value.id == "batch"
+                and target.attr == "input_ids"
+                and isinstance(node.value, ast.Call)
+                and isinstance(node.value.func, ast.Attribute)
+                and node.value.func.attr == "to"
+                and isinstance(node.value.func.value, ast.Name)
+                and node.value.func.value.id == "predict"
+            ):
+                assert any(
+                    isinstance(arg, ast.Attribute)
+                    and isinstance(arg.value, ast.Name)
+                    and arg.value.id == "torch"
+                    and arg.attr == "int64"
+                    for arg in node.value.args
+                )
+                return
+
+    raise AssertionError("draft extend input_ids must assign predict.to(torch.int64)")
+
+
+def test_eagle_draft_cuda_graph_padded_replay_updates_and_restores_seq_lens_sum():
+    """Padded CUDA graph replay must keep seq_lens_sum consistent.
+
+    Draft attention backends size/slice page tables from `seq_lens_sum`.  When
+    replay pads raw_bs to a captured bs, the metadata/replay path must see the
+    padded sum and the caller must get the raw sum restored afterwards.
+    """
+
+    tree = _parse_module("python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py")
+    func = _find_function(tree, "replay")
+
+    assigns = []
+    for node in ast.walk(func):
+        if not isinstance(node, ast.Assign):
+            continue
+        for target in node.targets:
+            if (
+                isinstance(target, ast.Attribute)
+                and isinstance(target.value, ast.Name)
+                and target.value.id == "forward_batch"
+                and target.attr == "seq_lens_sum"
+            ):
+                assigns.append(node)
+
+    assert len(assigns) >= 2, (
+        "replay must assign padded forward_batch.seq_lens_sum before graph "
+        "replay and restore the raw value afterwards"
+    )
+
+    assert any(
+        isinstance(node.value, ast.Name) and node.value.id == "raw_seq_lens_sum"
+        for node in assigns
+    ), "replay must restore raw_seq_lens_sum after padded graph replay"
+
+    assert any(
+        isinstance(node.value, ast.BinOp)
+        and any(
+            isinstance(child, ast.Name) and child.id == "raw_seq_lens_sum"
+            for child in ast.walk(node.value)
+        )
+        for node in assigns
+    ), "replay must derive a padded seq_lens_sum from raw_seq_lens_sum"
+
+
 def test_eagle_v2_binds_draft_runner_to_draft_extend_attention_backend():
     tree = _parse_module("python/sglang/srt/speculative/eagle_worker_v2.py")
     func = _find_function(tree, "init_attention_backend")