[htzo]

Motivation

Speculative decoding on CPU with the intel_amx attention backend produces output the target model would never have generated, causing the accuracy to be significantly worsened.

Fixes regression caused by #24013

Modifications

forward_extend and forward_decode both resolved seq_lens from the per-forward
metadata, then overwrote it:

seq_lens, extend_seq_lens, extend_start_loc, tree_mask = self.extend_metadata
...
seq_lens = forward_batch.seq_lens          # discards the resolved value
if seq_lens.dtype != torch.int64:
    seq_lens = seq_lens.to(torch.int64)

In TARGET_VERIFY mode _build_extend_metadata sets
seq_lens = forward_batch.seq_lens + num_draft_tokens, because extend_attention_cpu
derives the committed prefix by subtraction (seq_len_prefix = seq_lens - extend_seq_lens).
Overwritten, the kernel saw a prefix short by num_draft_tokens: the target scored every
draft token against a hole in its own context, and greedy verify accepted tokens its own
argmax would not have produced. forward_decode likewise discarded the expanded
per-candidate lengths used when speculative_num_steps > 1.

This deletes the two re-reads and keeps the int64 cast. Non-speculative extend is unaffected: there
forward_batch.seq_lens already equals prefix + extend.

Accuracy Tests

GSM8K with greedy sampling:

Target	Algorithm	steps/topk/draft	Baseline	Spec, before	Spec, after
MiMo-7B-RL	NEXTN	1 / 1 / 2	0.792	0.010	0.792
Llama-2-7b-chat-hf	EAGLE	3 / 1 / 4	0.244	0.006	0.244

Speed Tests and Profiling

Not measured - it's a correctness change

Checklist

• Format your code according to the Format code with pre-commit.
• Add unit tests according to the Run and add unit tests.
• Update documentation according to Write documentations.
• Provide accuracy and speed benchmark results according to Test the accuracy and Benchmark the speed.
• Follow the SGLang code style guidance.

Review and Merge Process

1. Ping Merge Oncalls to start the process. See the PR Merge Process.
2. Get approvals from CODEOWNERS and other reviewers.
3. Trigger CI tests with comments or contact authorized users to do so.
   • Common commands include /tag-and-rerun-ci, /tag-run-ci-label, /rerun-failed-ci
4. After green CI and required approvals, ask Merge Oncalls or people with Write permission to merge the PR.

---

CI States

Latest PR Test (Base): 🚫 Run #32925792115
Latest PR Test (Extra): ✅ Run #32937057026
Latest PR Test (AMD ROCm 7.2): ❌ Run #32925792202

[htzo]

/tag-run-ci-label

[mingfeima]

/rerun-failed-ci

[mingfeima]

check ci again.

[htzo]

check ci again.

• PR Test (Xeon) / build-test (xeon-gnr, gnr, base-b-tp-test-cpu, gnr) (pull_request) failed due to MMLU corpus 403

• ROCm failed due to connection issue and possibly runner issue? (Subprocess scheduler_0 (pid=15581) crashed with exit code -6 - I don't think it's related to our PR)

• PR Test (Base) failed is also failing upstream; [Unified Cache]: add glm5.2 per commit ci #36281 is probably the guilty PR

[ekintel]

Follow-up in #35881 for two potential gaps left behind: a regression test
covering both sites, and dropping the now-stale
disabled="EAGLE3 numerical parity mismatches on CPU intel_amx" marker from
test_spec_eagle_parity_cpu — those mismatches were the forward_extend
symptom, so it should be safe to re-enable now.

The test is verified in both directions on a Xeon: 4 passed on main,
and 2 failed, 2 passed against 058f99c5^, with only the two speculative
cases failing. It stubs the kernel entry points, so it needs no model or server.

@mingfeima @htzo would either of you be willing to add the run-ci label
there? It has had no CI run yet. Same for #35878 if you have a moment.