EngineX/xc-llm-ascend
3
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

[Feature] Support to use fullgraph with eagle (#5118)

Browse Source 
### What this PR does / why we need it?
    
We support to use full graph with eagle. 

Change list:
1. Distinguish between processing graph_params and draft_graph_params in
attention_v1.
    2. Adapt the full-graph mode in eagle_proposer, include:
        1). If use full graph, make Fullgraph Wrapper when load model.
2). Build a new meatadata, set running mode in FULL and mark attention
update in dummy_run when in Fullgraph mode.
3). Fixed and fill any attn_metadata, such as
attn_metadata.slot_mapping.
        4). Add a descriptor.
        5). Set running mode and triggered update metadata.
3. Trans is_mtp_model to is_draft_model, and add the update of
workspace.

NOTE:
When set async_scheduling=True, the draft model will enforce execution
in eager mode.

### Does this PR introduce _any_ user-facing change?

### How was this patch tested?

- vLLM version: v0.12.0
- vLLM main:
ad32e3e19c

---------

Signed-off-by: anon189Ty <Stari_Falcon@outlook.com>
Co-authored-by: Yizhou Liu <liu_yizhou@outlook.com>
Co-authored-by: Yizhou <136800916+yiz-liu@users.noreply.github.com>
This commit is contained in:
anon189Ty
2025-12-29 09:54:51 +08:00
committed by GitHub
parent f81cf694b2
commit 3e67e8276c
11 changed files with 348 additions and 103 deletions
Download Patch File Download Diff File Expand all files Collapse all files

41
tests/ut/compilation/test_acl_graph.py
View File

@@ -27,9 +27,9 @@ from vllm_ascend.attention.attention_v1 import (AscendMetadata,
from vllm_ascend.attention.mla_v1 import (AscendMLADecodeMetadata,
AscendMLAMetadata)
from vllm_ascend.compilation.acl_graph import (
ACLGraphEntry, ACLGraphWrapper, get_graph_params, get_mtp_graph_params,
set_graph_params, set_mtp_graph_params, update_attn_dcp_pcp_params,
update_mla_attn_dcp_pcp_params, update_mtp_graph_params_workspaces)
ACLGraphEntry, ACLGraphWrapper, get_draft_graph_params, get_graph_params,
set_draft_graph_params, set_graph_params, update_attn_dcp_pcp_params,
update_draft_graph_params_workspaces, update_mla_attn_dcp_pcp_params)
class TestACLGraphEntry(TestBase):
@@ -713,25 +713,26 @@ class TestACLGraphWrapper(TestBase):
self.assertEqual(unwrapped, self.mock_runnable)
class TestMTPGraphParams(TestBase):
class TestDraftGraphParams(TestBase):
def test_set_mtp_graph_params(self):
with patch('vllm_ascend.compilation.acl_graph._mtp_graph_params',
def test_set_draft_graph_params(self):
with patch('vllm_ascend.compilation.acl_graph._draft_graph_params',
new=None):
set_mtp_graph_params([4])
from vllm_ascend.compilation.acl_graph import _mtp_graph_params
self.assertIsNotNone(_mtp_graph_params)
set_draft_graph_params([4])
from vllm_ascend.compilation.acl_graph import _draft_graph_params
self.assertIsNotNone(_draft_graph_params)
@patch('vllm_ascend.compilation.acl_graph._mtp_graph_params')
def test_update_mtp_graph_params_workspaces(self, mtp_graph_params_mock):
mtp_graph_params_mock.workspaces = {4: 5}
update_mtp_graph_params_workspaces(4, 6)
self.assertEqual(mtp_graph_params_mock.workspaces[4], 6)
@patch('vllm_ascend.compilation.acl_graph._draft_graph_params')
def test_update_draft_graph_params_workspaces(self,
draft_graph_params_mock):
draft_graph_params_mock.workspaces = {4: 5}
update_draft_graph_params_workspaces(4, 6)
self.assertEqual(draft_graph_params_mock.workspaces[4], 6)
@patch('vllm_ascend.compilation.acl_graph._mtp_graph_params')
def test_get_mtp_graph_params(self, mtp_graph_params_mock):
graph_params = get_mtp_graph_params()
self.assertIs(mtp_graph_params_mock, graph_params)
@patch('vllm_ascend.compilation.acl_graph._draft_graph_params')
def test_get_draft_graph_params(self, draft_graph_params_mock):
graph_params = get_draft_graph_params()
self.assertIs(draft_graph_params_mock, graph_params)
class TestPCPDCPGraphParams(TestBase):
@@ -783,7 +784,7 @@ class TestPCPDCPGraphParams(TestBase):
decode=decode)
forward_context = MagicMock()
forward_context.attn_metadata = {"attn_layer_0": metadata}
forward_context.is_mtp_model = False
forward_context.is_draft_model = False
num_heads = 256
scale = 0.1
@@ -836,7 +837,7 @@ class TestPCPDCPGraphParams(TestBase):
decode_meta=decode)
forward_context = MagicMock()
forward_context.attn_metadata = {"attn_layer_0": metadata}
forward_context.is_mtp_model = False
forward_context.is_draft_model = False
self.graph_params.attn_params[4] = []
self.graph_params.attn_params[4].append(

73
tests/ut/spec_decode/test_eagle_proposer.py
View File

@@ -2,7 +2,7 @@ from unittest.mock import MagicMock, patch
import numpy as np
import torch
from vllm.config import CacheConfig, CompilationMode, VllmConfig
from vllm.config import CacheConfig, CompilationMode, CUDAGraphMode, VllmConfig
from tests.ut.base import TestBase
from vllm_ascend.spec_decode.eagle_proposer import EagleProposer
@@ -33,11 +33,13 @@ class TestEagleProposerInitialization(TestBase):
def tearDown(self):
self.mock_cpugpubuffer.stop()
def test_initialization_eagle(self):
def test_initialization_eagle_graph(self):
self.vllm_config.speculative_config.method = "eagle"
self.vllm_config.speculative_config.draft_model_config.get_hidden_size.return_value = 4096
self.vllm_config.compilation_config.mode = CompilationMode.VLLM_COMPILE
self.vllm_config.model_config.enforce_eager = False
self.vllm_config.speculative_config.enforce_eager = False
self.vllm_config.scheduler_config.async_scheduling = False
proposer = EagleProposer(vllm_config=self.vllm_config,
device=self.device,
@@ -53,7 +55,7 @@ class TestEagleProposerInitialization(TestBase):
self.assertEqual(proposer.hidden_states.shape, (1024, 4096))
self.assertEqual(proposer.arange.shape, (1024, ))
def test_initialization_eagle3(self):
def test_initialization_eagle3_enforce_eager(self):
self.vllm_config.speculative_config.method = "eagle3"
self.vllm_config.speculative_config.draft_model_config.get_hidden_size.return_value = 2048
self.vllm_config.compilation_config.mode = CompilationMode.NONE
@@ -68,6 +70,23 @@ class TestEagleProposerInitialization(TestBase):
self.assertFalse(proposer.use_cuda_graph)
self.assertEqual(proposer.hidden_states.shape, (1024, 2048))
def test_initialization_eagle3_full_graph_async(self):
self.vllm_config.speculative_config.method = "eagle3"
self.vllm_config.speculative_config.draft_model_config.get_hidden_size.return_value = 2048
self.vllm_config.compilation_config.mode = CompilationMode.VLLM_COMPILE
self.vllm_config.model_config.enforce_eager = False
self.vllm_config.speculative_config.enforce_eager = False
self.vllm_config.scheduler_config.async_scheduling = True
proposer = EagleProposer(vllm_config=self.vllm_config,
device=self.device,
runner=self.runner)
self.assertEqual(proposer.name, SpecDcodeType.EAGLE3)
self.assertEqual(proposer.hidden_size, 2048)
self.assertFalse(proposer.use_cuda_graph)
self.assertEqual(proposer.hidden_states.shape, (1024, 2048))
class TestEagleProposerLoadModel(TestBase):
@@ -176,6 +195,7 @@ class TestEagleProposerDummyRun(TestBase):
def setUp(self):
self.vllm_config = MagicMock(spec=VllmConfig)
self.vllm_config.speculative_config = MagicMock()
self.vllm_config.speculative_config.num_speculative_tokens = 4
self.device = torch.device("cpu")
self.runner = MagicMock()
@@ -192,25 +212,64 @@ class TestEagleProposerDummyRun(TestBase):
device=self.device,
runner=self.runner)
self.proposer.model = MagicMock()
self.proposer.update_stream = MagicMock()
def tearDown(self):
self.mock_cpugpubuffer.stop()
@patch("vllm_ascend.spec_decode.eagle_proposer.get_forward_context")
@patch("vllm_ascend.spec_decode.eagle_proposer.set_ascend_forward_context")
def test_dummy_run_basic(self, mock_context):
def test_dummy_run_basic(self, mock_context, mock_get_context):
num_tokens = 32
with_prefill = False
self.proposer.dummy_run(num_tokens=num_tokens,
with_prefill=with_prefill)
mock_context.assert_called_once()
self.assertTrue(self.proposer.model.call_count == 4)
@patch("vllm_ascend.spec_decode.eagle_proposer.get_forward_context")
@patch("vllm_ascend.spec_decode.eagle_proposer.set_ascend_forward_context")
def test_dummy_run_with_prefill(self, mock_context):
def test_dummy_run_with_prefill(self, mock_context, mock_get_context):
mock_context.return_value.__enter__.return_value = None
self.proposer.dummy_run(num_tokens=64, with_prefill=True, num_reqs=4)
self.proposer.model.assert_called_once()
self.assertTrue(self.proposer.model.call_count == 4)
@patch("vllm_ascend.spec_decode.eagle_proposer.update_attn_params")
@patch("vllm_ascend.spec_decode.eagle_proposer.get_forward_context")
@patch("vllm_ascend.spec_decode.eagle_proposer.set_ascend_forward_context")
def test_dummy_run_in_graph_capture(self, mock_context, mock_get_context,
mock_update_attn_params):
last_use_cuda_graph = self.proposer.use_cuda_graph
mock_return_context = MagicMock()
mock_return_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
mock_return_context.capturing = True
mock_get_context.return_value = mock_return_context
self.proposer.use_cuda_graph = True
self.proposer.dummy_run(num_tokens=64,
in_graph_capturing=True,
aclgraph_runtime_mode=CUDAGraphMode.FULL)
self.assertTrue(self.proposer.model.call_count == 4)
mock_update_attn_params.assert_not_called()
self.proposer.use_cuda_graph = last_use_cuda_graph
@patch("vllm_ascend.spec_decode.eagle_proposer.update_attn_params")
@patch("vllm_ascend.spec_decode.eagle_proposer.get_forward_context")
@patch("vllm_ascend.spec_decode.eagle_proposer.set_ascend_forward_context")
def test_dummy_run_in_graph_run(self, mock_context, mock_get_context,
mock_update_attn_params):
last_use_cuda_graph = self.proposer.use_cuda_graph
mock_return_context = MagicMock()
mock_return_context.cudagraph_runtime_mode = CUDAGraphMode.FULL
mock_return_context.capturing = False
mock_get_context.return_value = mock_return_context
self.proposer.use_cuda_graph = True
self.proposer.dummy_run(num_tokens=64,
in_graph_capturing=False,
aclgraph_runtime_mode=CUDAGraphMode.FULL)
self.assertTrue(self.proposer.model.call_count == 4)
self.assertTrue(mock_update_attn_params.call_count == 4)
self.proposer.use_cuda_graph = last_use_cuda_graph
class TestEagleProposerGenerateTokenIds(TestBase):