From c6b99f6060d7f3931789eb1b8c54b7eaff908e18 Mon Sep 17 00:00:00 2001 From: laoyao0822 Date: Sat, 27 Jun 2026 02:35:09 +0800 Subject: [PATCH] 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 --- .../eagle_draft_cuda_graph_runner.py | 10 ++- .../sglang/srt/speculative/eagle_info_v2.py | 5 +- .../test_eagle_v2_draft_extend_contract.py | 82 +++++++++++++++++++ 3 files changed, 95 insertions(+), 2 deletions(-) diff --git a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py b/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py index f80a38a64..fccb286e7 100644 --- a/python/sglang/srt/speculative/eagle_draft_cuda_graph_runner.py +++ b/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 diff --git a/python/sglang/srt/speculative/eagle_info_v2.py b/python/sglang/srt/speculative/eagle_info_v2.py index a095bf988..b5d479f2b 100644 --- a/python/sglang/srt/speculative/eagle_info_v2.py +++ b/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 diff --git a/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py b/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py index 84f5fe2bc..70e5a9a95 100644 --- a/test/registered/spec/eagle/test_eagle_v2_draft_extend_contract.py +++ b/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")