Skip to content

[Bugfix][Spec Decode] Validate YaRN for extended native MTP context - #2

Merged
ga-it merged 1 commit into
mainfrom
review/native-mtp-yarn-context
Aug 28, 2026
Merged

ga-it merged 1 commit into
mainfrom
review/native-mtp-yarn-context

Conversation

@ga-it

@ga-it ga-it commented Aug 28, 2026

Copy link
Copy Markdown
Owner

[Bugfix][Spec Decode] Validate YaRN for extended native MTP context

AI-assisted draft. AI assistance was used for upstream comparison,
implementation, tests, live profiling, and this write-up. Upstream requires
the human submitter to review every changed line and run the relevant tests
before submitting this PR. This note must not be removed; once completed, the
submitter should add their human-review confirmation and final Linux test
results.

Purpose

A same-checkpoint native MTP drafter can retain the checkpoint's native
max_position_embeddings even when the target is explicitly extended with
YaRN. For Qwen3.8 this left the drafter at 262,144 while the target was served
at a larger context. Above the drafter limit, speculative decoding fell back to
target-only decoding and the changing max_seq_len argument defeated
warning_once, producing one warning per output step.

This PR:

  • extends only same-checkpoint native MTP drafters;
  • requires explicit target rope_parameters with rope_type: yarn;
  • validates finite numeric values, factor > 1, and an
    original_max_position_embeddings that exactly matches the drafter's native
    limit;
  • refuses a requested limit above either the target limit or the validated
    original * factor YaRN ceiling;
  • deep-copies the validated target parameters to the drafter; and
  • makes the over-limit fallback warning stable and globally one-shot.

It does not invent RoPE scaling, change the target limit, change unrelated
draft models, or silently extend an unvalidated drafter.

Base SHA: 62ad1e02693f4c857f3b7547cef1860ee54e8053.

Context relationship

The validated V100 checkpoint has a native target/drafter limit of 262,144.
Factor-4 YaRN gives a mathematical ceiling of 1,048,576. The tested service
caps both target and native drafter at 1,000,000, leaving 48,576 tokens of
margin. The former 320,000 and 512,000 values were rollout tiers, not model
limits.

Duplicate-work audit

The complete open PR list was refreshed on 2026-08-28 at the base SHA; no open
PR contains this fail-closed inheritance or stable fallback warning.

  • PR 241 is an old-base,
    rejected aggregate. It unconditionally overwrites the
    drafter length and changes automatic NTK behavior. This PR instead accepts
    only an explicit, internally consistent target YaRN configuration and
    otherwise fails closed.
  • Merged PR 389,
    PR 391, and
    PR 399, plus open
    PR 398, improve V100 MTP4
    execution but do not extend the drafter's validated context.
  • PR 408 addresses Qwen3.8
    runtime correctness, not YaRN context inheritance.
  • Issue 214 and
    issue 93 describe separate
    long-prefill and non-MTP paths and are not claimed as fixed here.

Test Plan

Local current-main checks:

uvx --from ruff==0.14.0 ruff format --check \
  vllm/config/speculative.py \
  vllm/v1/worker/gpu_model_runner.py \
  tests/v1/spec_decode/test_max_len.py
3 files already formatted

uvx --from ruff==0.14.0 ruff check <same files>
All checks passed!

uv run --no-project --python 3.12 python -m py_compile <same files>
exit 0

git diff --check
exit 0

git apply --check 0001-upstream-native-mtp-yarn-context.patch
exit 0 against 62ad1e0

Focused CPU tests were added for:

  • valid YaRN inheritance and deep-copy isolation;
  • missing, dynamic, non-numeric, and native-limit-mismatched parameters;
  • rejection above the 1,048,576 validated ceiling; and
  • no mutation when the drafter is not extended.

Test Result

The same formatting, lint, byte-compilation, and diff checks also pass in a
disposable Linux worktree on gazasrv16, with the patch applied to the stated
base SHA. No production container or GPU was touched.

A focused pytest run was attempted there. An isolated editable environment
from unmodified current main first failed dependency resolution because of the
separate TokenSpeed/TVM-FFI conflict described in the companion candidate. A
test-only overlay of that companion patch allowed the precompiled editable
install (191 packages) and the full CUDA test dependency set to install. Pytest
then stopped during collection because the precompiled editable source layout
did not provide vllm._C. A CPU-source fallback compiled most targets but was
blocked by the host's missing numa.h development header and a Torch CPU API
mismatch. The production host packages were deliberately not changed.

Accordingly, the new test cases are present but no pytest pass is claimed. The
human submitter must run them in the prescribed, fully provisioned Linux UV
environment before upstream submission:

.venv/bin/python -m pytest tests/v1/spec_decode/test_max_len.py -v
pre-commit run --files \
  vllm/config/speculative.py \
  vllm/v1/worker/gpu_model_runner.py \
  tests/v1/spec_decode/test_max_len.py

Live V100 validation

The exact source behavior was also built and exercised in a running SM70
image. These measurements are specifically for four V100-SXM2-32GB GPUs with
TP4; they are not Blackwell results.

Setting Value
Model philbert440/Qwen3.8-27B-W4A16-AWQ
Hardware 4 x V100-SXM2-32GB, tensor parallel 4
Image 1cat-vllm-sm70:qwen38-27b-awq-main-62ad1e0-mtp4-yarn-pr331-cu128
Image digest sha256:d0fdeefbea5b61a12caa75e57543e06a01b2dadf3ea4eb74d0dfa0a90a48b95a
CUDA / KV CUDA 12.8.1; FP8 E5M2 KV cache
Speculation Native MTP4; target and drafter both 1,000,000
KV capacity 2,557,299 tokens, or 2.56 x 1M sequences; max-num-seqs=2

HTTP /health, /v1/models, and /metrics all returned 200. The service
remained healthy with zero restarts. All four V100s saturated during cold
prefill; peak observed KV occupancy was 35.1%.

Performance

Route / request TTFT Decode Acceptance Total
Current MTP4, 6,316 prompt / 512 output 2.9538 s 87.0357 tok/s 73.2824% 8.8249 s
Current MTP4, 564,577 cold / 512 output 690.9660 s 31.2186 tok/s 79.7131% 707.3345 s
Current MTP4, 564,580 cold with changed salt / 512 output 690.7554 s 32.0822 tok/s 82.3529% 706.6832 s
Current MTP4, identical 564,577 replay / 512 output 9.2287 s 31.1965 tok/s 79.7131% 25.6088 s
Earlier target-only above drafter cap, 564,575 / 512 545.395 s 8.772 tok/s none 603.647 s
Earlier validated MTP2, 564,575 / 512 552.084 s 22.874 tok/s 81.70% 574.423 s

At this long prompt, current MTP4 decode was 3.56x-3.66x the target-only
baseline and 1.37x-1.40x the earlier MTP2 result. The total-latency crossover
was approximately 1,800 generated tokens versus target-only and 11,000-12,000
versus MTP2 because current-main cold prefill is slower.

The current-main cold TTFT of about 691 seconds is roughly 139-146 seconds
slower than the earlier target-only/MTP2 baselines. This PR does not claim to
fix that separate prefill regression.

Tokenization and prefix cache

Backend tokenization took 2.4811-2.51695 seconds for roughly 564.6k tokens,
equivalent to 224,311-227,549 tokens/s, so tokenization was not the cold-query
bottleneck. The identical replay reused 562,368 of 564,577 queried prefix
tokens (99.6087%), reducing TTFT from about 691 seconds to 9.2287 seconds.

Warning and fallback behavior

  • Before this change, one OpenWebUI request ended at 326,399 sequence tokens,
    generated 18,167 output tokens, and emitted 18,169 changing-argument
    over-limit warnings.
  • With the stable fallback warning, a 321,914-token request above a staged 320k
    drafter tier generated 128 tokens with zero drafts and exactly one warning.
  • A later 520,339-token request above a staged 512k tier also cleanly used zero
    drafts and emitted one warning.
  • The 564.6k MTP4 runs under the 1M validated limit emitted no MTP warning
    flood. Each cold request emitted one separate Transformers advisory because
    the original checkpoint metadata still declares 262,144.

Real OpenWebUI workload

A production OpenWebUI agent request on the target-only control received
733,366 prompt tokens, reused 720,000 prefix tokens, and generated 23,724
tokens at approximately 8.5-9.2 tok/s. It completed without a proxy timeout,
queue, engine error, or restart. Its long generation lies well beyond the
measured MTP4 total-latency crossover and motivated keeping the MTP4 route for
Hermes/OpenClaw workloads.

Validation boundary

The service booted with matching 1,000,000 target/drafter limits and MTP was
tested end to end through 564,580 prompt tokens. A near-1M cold prefill has not
yet been executed, so this PR does not claim an end-to-end measurement at the
absolute configured limit.

Risk and rollback

Invalid or implicit scaling fails during configuration instead of failing
mid-request. Unextended native MTP and separate draft checkpoints are unchanged.
At runtime, requests above an intentionally lower drafter cap still fall back
to target-only decoding; only repetitive logging is suppressed. Rollback is a
revert of these three files.


Essential Elements of an Effective PR Description Checklist
  • The purpose of the PR, including the related PRs and issues that were checked without incorrectly claiming they are resolved.
  • The test plan, including the focused pytest and pre-commit commands required before upstream submission.
  • The test results, including Linux source checks, test-environment limitations, SM70/V100 E2E results, decode performance, MTP acceptance, TTFT, KV-cache capacity/occupancy, prefix-cache reuse, tokenization rate, health, and restart status.
  • (Optional) The necessary documentation update was considered. No public API, CLI, or supported-model surface changes; focused tests describe the validation contract.

BEFORE SUBMITTING, PLEASE READ https://docs.vllm.ai/en/latest/contributing (anything written below this line will be removed by GitHub Actions)

@ga-it
ga-it marked this pull request as ready for review August 28, 2026 18:41
@ga-it
ga-it merged commit 4bcc1e1 into main Aug 28, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant