Skip to content

Commit f6a737c

Browse files
[Feature] Virtual Width Network Eagle3 (vllm-project#10042)
### What this PR does / why we need it? We add a customed speculative model "Virtual Width Network Eagle3" with architecture similar to Seed's work 'Virtual Width Network' (https://huggingface.co/papers/2511.11238) We impelmented and trained the 'vwn eagle3' based on 'speculators' library (https://github.com/vllm-project/speculators); the training datasets includes 'ultrachat' (https://huggingface.co/datasets/openbmb/UltraChat) and 'sharegpt' (https://huggingface.co/collections/bunnycore/sharegpt-datasets). We used the same setting to train a vwn-eagle3 and an eagle3 model to compare their performace on longbenchv2, gsm8k, sharegpt, and humaneval; as well as some man-made fixed-length datasets. The performance is concluded as following: Datasets testing result: | Dataset | Context-Length | E2E Throughput Gain| | ------ | ------ | ------ | | gsm8k | 1.5k | 7.9%| | humaneval | 2.5k | 11.94%| | sharegpt | 5k | 2.84%| | longbench_v2| 128k | 18.23%| Fixed length testing results: |Input Lenght | Output Length | E2E Throughput Gain| | ------ | ------ | ------ | | 2k | 2k | 9.83%| | 3.5k | 1.5k | 2.38%| | 32k | 1k | 18.66%| ### Does this PR introduce _any_ user-facing change? Nope. VWN Eagle3 uses the same interface with eagle3 and can be distinguished by draft model's config file, thus, the user can simply use the same setting to start the server: `--speculative-config '{"method": "eagle3","model": "/vwn-eagle3-model-path/", "num_speculative_tokens": 3}' ` To reduce redundant codes, we reuse the implementation of llama_eagle3.py to the greatest extent: | Component | Provided by Upstream Eagle3 | Overridden/Added by VWN | Reuse Rate | |-----------|-----------------------------|------------------------|------------| | Top-level CausalLM | ✅ forward, embed_input_ids, load_weights | Only replaces `self.model` | ~80% | | Model container | ✅ embed, fc, load_weights | Swaps layer type + simplifies forward | ~50% | | Decoder Layer | ✅ self_attn, mlp, quant | Adds VWN projections + rewrites forward | ~40% | | PreVwnLayerV1 | ❌ No counterpart | Entirely new implementation | 0% (pure addition) | ### How was this patch tested? - vLLM-ascend main @ 7cb4484 - vLLM version: v0.22.1 - vLLM main: vllm-project/vllm@967c5c3 --------- Signed-off-by: YzTongNiar <1667927948@qq.com> Co-authored-by: wu_yushan <wuyushan1@huawei.com>
1 parent 16e9156 commit f6a737c

7 files changed

Lines changed: 796 additions & 2 deletions

File tree

.github/workflows/misc/model_dataset_list.json

Lines changed: 2 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -274,7 +274,8 @@
274274
"Eco-Tech/GLM-5.1-w8a8",
275275
"Eco-Tech/Qwen3.5-397B-A17B-w4a8-mtp",
276276
"lightseekorg/kimi-k2.5-eagle3",
277-
"Eco-Tech/Qwen3.5-122B-A10B-w8a8-mtp"
277+
"Eco-Tech/Qwen3.5-122B-A10B-w8a8-mtp",
278+
"vllm-ascend/Qwen3-30B-A3B-vwn-eagle-model"
278279
],
279280
"datasets": [
280281
"vllm-ascend/GSM8K-in1024-bs210",

.github/workflows/scripts/test_config.yaml

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -451,11 +451,13 @@
451451
source_file_dependencies:
452452
- vllm_ascend/spec_decode/llm_base_proposer.py
453453
- vllm_ascend/spec_decode/eagle_proposer.py
454+
- vllm_ascend/models/llama_eagle3_vwn.py
454455
tests:
455456
- tests/e2e/pull_request/one_card/spec_decode/test_eagle.py
456457
- tests/e2e/pull_request/one_card/spec_decode/test_mtp_eagle_correctness.py
457458
- tests/e2e/pull_request/two_card/spec_decode/test_spec_decode.py
458459
- tests/e2e/pull_request/four_card/spec_decode/test_mtp_qwen3_next.py
460+
- tests/ut/spec_decode/test_speculators_vwn_eagle3.py
459461

460462
- name: spec_decode_ngram
461463
optional: false

tests/e2e/pull_request/two_card/spec_decode/test_spec_decode.py

Lines changed: 100 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,11 +46,19 @@
4646
},
4747
}
4848

49+
VWN_EAGLE3_MODELS = {
50+
"vwn_eagle3": {
51+
"main": "Qwen/Qwen3-30B-A3B",
52+
"spec": "vllm-ascend/Qwen3-30B-A3B-vwn-eagle-model",
53+
},
54+
}
55+
4956
# NOTE: golden may change (eagle_proposer only runs in eager mode currently),
5057
# thus please update it if ci fails but you have better acceptance
5158
BASELINES_SP = {
5259
"eagle3": [0.68, 0.40, 0.18],
5360
"p-eagle": [0.5625, 0.25, 0.0625, 0.0, 0.0, 0.0, 0.0, 0.0],
61+
"vwn_eagle3": [0.75, 0.5, 0.3],
5462
}
5563

5664

@@ -327,3 +335,95 @@ def test_p_eagle_acceptance(
327335
print(f"golden: {golden}")
328336

329337
assert match
338+
339+
340+
@patch.dict(os.environ, {"VLLM_ASCEND_ENABLE_FLASHCOMM1": "1"})
341+
def test_qwen3_vwn_eagle3_tp2():
342+
"""
343+
Test Qwen3-30B-A3B with VWN-Eagle3 speculative decoding acceptance rate.
344+
This test verifies that VWN-Eagle3 spec decode works correctly with:
345+
- Tensor Parallel size = 4
346+
- Expert Parallel enabled (for MoE)
347+
- num_speculative_tokens = 3
348+
- enforce_eager = True
349+
- Acceptance rate matches baseline (tolerance 0.06)
350+
"""
351+
num_speculative_tokens = 3
352+
main_model_name = VWN_EAGLE3_MODELS["vwn_eagle3"]["main"]
353+
spec_model_name = VWN_EAGLE3_MODELS["vwn_eagle3"]["spec"]
354+
355+
tokenizer = AutoTokenizer.from_pretrained(
356+
main_model_name,
357+
trust_remote_code=True,
358+
)
359+
sampling_params = SamplingParams(
360+
temperature=0,
361+
ignore_eos=False,
362+
max_tokens=256,
363+
)
364+
365+
prompts = [
366+
{
367+
"role": "user",
368+
"content": "Hello, my name is",
369+
},
370+
{
371+
"role": "user",
372+
"content": "The capital of France is",
373+
},
374+
{
375+
"role": "user",
376+
"content": "The future of AI is",
377+
},
378+
]
379+
prompts = [
380+
tokenizer.apply_chat_template(
381+
[prompt],
382+
tokenize=False,
383+
add_generation_prompt=True,
384+
)
385+
for prompt in prompts
386+
]
387+
388+
speculative_config = {
389+
"method": "eagle3",
390+
"num_speculative_tokens": num_speculative_tokens,
391+
"model": spec_model_name,
392+
}
393+
394+
with VllmRunner(
395+
main_model_name,
396+
enforce_eager=True,
397+
max_model_len=2048,
398+
disable_log_stats=False,
399+
tensor_parallel_size=2,
400+
max_num_seqs=16,
401+
distributed_executor_backend="mp",
402+
gpu_memory_utilization=0.92,
403+
speculative_config=speculative_config,
404+
enable_expert_parallel=True,
405+
) as llm:
406+
_ = llm.generate(prompts, sampling_params)
407+
metrics = llm.model.get_metrics()
408+
409+
# Check acceptance rate
410+
num_drafts = 0
411+
num_accepted_tokens_per_pos = [0] * num_speculative_tokens
412+
for metric in metrics:
413+
if metric.name == "vllm:spec_decode_num_drafts":
414+
assert isinstance(metric, Counter)
415+
num_drafts += metric.value
416+
elif metric.name == "vllm:spec_decode_num_accepted_tokens_per_pos":
417+
assert isinstance(metric, Vector)
418+
for pos in range(len(metric.values)):
419+
num_accepted_tokens_per_pos[pos] += metric.values[pos]
420+
421+
acceptance_per_pos = [n / num_drafts for n in num_accepted_tokens_per_pos]
422+
golden = BASELINES_SP["vwn_eagle3"]
423+
424+
match = all(abs(a - b) < 0.06 for a, b in zip(acceptance_per_pos, golden))
425+
if not match:
426+
print(f"acceptance_per_pos: {acceptance_per_pos}")
427+
print(f"golden: {golden}")
428+
429+
assert match

0 commit comments

Comments
 (0)