[GLM5.3 Perf] Optimize glm 5.3 metadata op, 1.6~4.8x kernel level performance improvement (#58450)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 25, 2026 at 13:15 UTC
c88f4824fff42794d026348b86f9b78869fa273e
4 files changed
+26
-25
tests/v1/attention/test_indexer_deepseek_v4_slot_mapping.py
+19
-4
@@ -7,6 +7,9 @@ import pytest
7
import torch
8
9
from tests.v1.attention.utils import create_vllm_config
10
+from vllm.model_executor.layers.attention.sparse_mla_attention import (
11
+ SparseMLACommonMetadataBuilder,
12
+)
13
from vllm.models.deepseek_v4.sparse_mla import DeepseekV4SparseMLABackend
14
from vllm.models.deepseek_v41.sparse_mla import (
15
DeepseekV4SparseMLABackend as DeepseekV41SparseMLABackend,
@@ -95,8 +98,11 @@ def test_fused_indexer_decode_metadata(query_lens, padding):
98
99
100
@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires CUDA")
98
-@pytest.mark.parametrize("query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0]])
99
-def test_device_token_request_mapping(query_lens):
101
+@pytest.mark.parametrize(
102
+ "query_lens", [[1], [0, 257, 1, 0, 3], [6] * 64, [2, 6, 4, 0], [0, 0, 0]]
103
+)
104
+@pytest.mark.parametrize("use_sparse_mla_builder", [False, True])
105
+def test_device_token_request_mapping(query_lens, use_sparse_mla_builder):
106
"""Graph replay follows device boundaries even when CPU lengths are stale."""
107
lengths = torch.tensor(query_lens, device="cuda", dtype=torch.int32)
108
qsl = torch.cat(
@@ -117,7 +123,16 @@ def test_device_token_request_mapping(query_lens):
123
),
124
slot_mapping=torch.full((n + 7,), -1, device="cuda", dtype=torch.int64),
125
)
120
- result = common.token_to_req_indices(output)
126
+
127
+ def build_mapping():
128
+ if use_sparse_mla_builder:
129
+ builder = SimpleNamespace(req_id_per_token_buffer=output)
130
+ return SparseMLACommonMetadataBuilder._build_req_id_per_token(
131
+ builder, common
132
+ )
133
+ return common.token_to_req_indices(output)
134
+
135
+ result = build_mapping()
136
assert result.data_ptr() == output.data_ptr()
137
expected = torch.repeat_interleave(
138
torch.arange(len(query_lens), device="cuda", dtype=torch.int32), lengths
@@ -127,7 +142,7 @@ def test_device_token_request_mapping(query_lens):
142
graph = torch.cuda.CUDAGraph()
143
common._token_to_req_indices_cache = None
144
with torch.cuda.graph(graph):
130
- common.token_to_req_indices(output)
145
+ build_mapping()
146
reversed_lens = lengths.flip(0)
147
qsl[1:].copy_(reversed_lens.cumsum(0))
148
graph.replay()
vllm/model_executor/layers/attention/sparse_mla_attention.py
+1
-11
@@ -291,17 +291,7 @@ class SparseMLACommonMetadataBuilder(AttentionMetadataBuilder[T]):
291
self,
292
common_attn_metadata: "CommonAttentionMetadata",
293
) -> torch.Tensor:
294
- num_tokens = common_attn_metadata.num_actual_tokens
295
- starts = np.asarray(common_attn_metadata.query_start_loc_cpu, dtype=np.int32)
296
- seg_lengths = np.diff(starts)
297
- req_id_per_token = np.repeat(
298
- np.arange(seg_lengths.shape[0], dtype=np.int32), seg_lengths
299
- )
300
- self.req_id_per_token_buffer.fill_(0)
301
- self.req_id_per_token_buffer[: req_id_per_token.shape[0]].copy_(
302
- np_to_pinned_tensor(req_id_per_token), non_blocking=True
303
- )
304
- return self.req_id_per_token_buffer[:num_tokens]
294
+ return common_attn_metadata.token_to_req_indices(self.req_id_per_token_buffer)
295
296
def _build_chunked_context_fields(
297
self,
vllm/v1/attention/backends/mla/flashattn_mla_sparse.py
+6
-3
@@ -178,6 +178,11 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
178
assert self.topk_indices_buffer is not None, (
179
"Indexer or topk_indices_buffer required for sparse MLA"
180
)
181
+ self.cu_seqlens_q_buffer = torch.arange(
182
+ self.topk_indices_buffer.shape[0] + 1,
183
+ dtype=torch.int32,
184
+ device=self.topk_indices_buffer.device,
185
+ )
186
self.supports_quant_query_input = False
187
188
def forward_mqa(
@@ -306,9 +311,7 @@ class FlashAttnMLASparseImpl(SparseMLACommonImpl[FlashAttnMLASparseMetadata]):
311
kv_cache if cache_is_flat else flat_kv_row_view(kv_cache, block_size)[0]
312
)
313
309
- cu_seqlens_q = torch.arange(
310
- 0, q_rope.shape[0] + 1, dtype=torch.int32, device=q_rope.device
311
- )
314
+ cu_seqlens_q = self.cu_seqlens_q_buffer[: q_rope.shape[0] + 1]
315
v_cache = kv_rows[:, : self.kv_lora_rank].unsqueeze(1).unsqueeze(1)
316
if self.qk_rope_head_dim == 0:
317
# FA3's QV path requires the 64-wide Q/K specialization. For NoPE
vllm/v1/attention/backends/mla/flashinfer_mla_sparse.py
-7
@@ -38,7 +38,6 @@ from vllm.v1.kv_cache_interface import AttentionSpec
38
39
if TYPE_CHECKING:
40
from vllm.model_executor.models.deepseek_v2 import Indexer
41
- from vllm.v1.attention.backend import CommonAttentionMetadata
41
42
logger = init_logger(__name__)
43
@@ -301,12 +300,6 @@ class FlashInferMLASparseTRTLLMMetadataBuilder(FlashInferMLASparseMetadataBuilde
300
),
301
)
302
304
- def _build_req_id_per_token(
305
- self,
306
- common_attn_metadata: "CommonAttentionMetadata",
307
- ) -> torch.Tensor:
308
- return common_attn_metadata.token_to_req_indices(self.req_id_per_token_buffer)
309
-
303
304
# Global workspace buffer (lazily initialized)
305
_fi_sparse_workspace: torch.Tensor | None = None