From 54d9b1cb487c2a2f12506942820f0ae443807093 Mon Sep 17 00:00:00 2001 From: Wentao Ye <44945378+yewentao256@users.noreply.github.com> Date: Thu, 15 Oct 2026 16:20:12 -0400 Subject: [PATCH] [GLM5.3 Perf] Inplace update draft decode metadata, 1.7% E2E Throughput Improvement (#61873) Signed-off-by: yewentao256 --- .../backends/mla/flashinfer_mla_sparse.py | 18 ++++++++++++++++++ 1 file changed, 18 insertions(+) diff --git a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py index d673868975..605f551bc2 100644 --- a/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py +++ b/vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py @@ -35,6 +35,7 @@ from vllm.v1.attention.backends.mla.sparse_utils import ( flat_kv_row_view, prepare_sparse_mla_safe_lengths, ) +from vllm.v1.attention.ops.metadata import compute_token_to_req_indices from vllm.v1.kv_cache_interface import AttentionSpec, KVCacheSpec if TYPE_CHECKING: @@ -297,6 +298,11 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde device: torch.device, ) -> None: super().__init__(kv_cache_spec, layer_names, vllm_config, device) + self.supports_draft_decode_metadata_update = ( + self.dcp_world_size == 1 + and self.pcp_world_size == 1 + and vllm_config.attention_config.hisparse_config is None + ) # Under DCP the workspace must hold the trtllm-gen softmax-stats slab. # The buffer address is baked into CUDA graphs, so it has to reach its @@ -314,6 +320,18 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde ), ) + def update_draft_decode_metadata( + self, metadata: FlashInferMLASparseMetadata + ) -> None: + num_tokens = metadata.num_decode_tokens + if num_tokens: + compute_token_to_req_indices( + metadata.query_start_loc, + metadata.req_id_per_token, + num_tokens, + num_tokens, + ) + # Global workspace buffer (lazily initialized) _fi_sparse_workspace: torch.Tensor | None = None