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(