[GLM5.3 Perf] Size the GLM-5 sparse indexer decode workspace, 3072 MiB GPU memory saved (#57701)
Signed-off-by: yewentao256 <zhyanwentao@126.com>
Wentao Ye committed
Sep 19, 2026 at 16:57 UTC
36fa72d2d0d2f86c7c83e1e99c9012b7bd26463b
3 files changed
+17
-17
vllm/models/glm5next/amd/sparse_indexer.py
+8
-8
@@ -100,7 +100,7 @@ def sparse_attn_indexer_kpool(
100
scale_fmt: str | None,
101
topk_tokens: int,
102
head_dim: int,
103
- max_model_len: int,
103
+ max_pool_len: int,
104
total_seq_lens: int,
105
topk_indices_buffer: torch.Tensor,
106
skip_k_cache_insert: bool,
@@ -136,7 +136,7 @@ def sparse_attn_indexer_kpool(
136
)
137
138
# Reserve profiler-visible memory for the worst-case decode logits,
139
- # whose shape is [B * next_n, max_model_len]. This profiling branch
139
+ # whose shape is [B * next_n, max_pool_len]. This profiling branch
140
# returns before invoking the logits kernel itself.
141
cfg = get_current_vllm_config_or_none()
142
worst_decode_tokens = 0
@@ -152,7 +152,7 @@ def sparse_attn_indexer_kpool(
152
sched.max_num_batched_tokens,
153
)
154
# float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
155
- decode_logits_elems = worst_decode_tokens * max_model_len * 4
155
+ decode_logits_elems = worst_decode_tokens * max_pool_len * 4
156
prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
157
max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
158
_ = torch.empty(
@@ -566,7 +566,7 @@ def sparse_attn_indexer_kpool(
566
seq_lens,
567
decode_metadata.block_table,
568
decode_metadata.schedule_metadata,
569
- max_model_len=max_model_len,
569
+ max_model_len=max_pool_len,
570
)
571
num_rows = logits.shape[0]
572
# kpool: logits are pool-granular -> select topk_tokens//kpool pools,
@@ -645,7 +645,7 @@ class SparseAttnIndexerKpool(CustomOp):
645
scale_fmt: str,
646
topk_tokens: int,
647
head_dim: int,
648
- max_model_len: int,
648
+ max_pool_len: int,
649
max_total_seq_len: int,
650
topk_indices_buffer: torch.Tensor,
651
skip_k_cache_insert: bool = False,
@@ -659,7 +659,7 @@ class SparseAttnIndexerKpool(CustomOp):
659
self.scale_fmt = scale_fmt
660
self.topk_tokens = topk_tokens
661
self.head_dim = head_dim
662
- self.max_model_len = max_model_len
662
+ self.max_pool_len = max_pool_len
663
self.max_total_seq_len = max_total_seq_len
664
self.topk_indices_buffer = topk_indices_buffer
665
self.skip_k_cache_insert = skip_k_cache_insert
@@ -698,7 +698,7 @@ class SparseAttnIndexerKpool(CustomOp):
698
self.scale_fmt,
699
self.topk_tokens,
700
self.head_dim,
701
- self.max_model_len,
701
+ self.max_pool_len,
702
self.max_total_seq_len,
703
self.topk_indices_buffer,
704
skip_k_cache_insert=self.skip_k_cache_insert,
@@ -715,7 +715,7 @@ class SparseAttnIndexerKpool(CustomOp):
715
self.scale_fmt,
716
self.topk_tokens,
717
self.head_dim,
718
- self.max_model_len,
718
+ self.max_pool_len,
719
self.max_total_seq_len,
720
self.topk_indices_buffer,
721
self.skip_k_cache_insert,
vllm/models/glm5next/common/attention.py
+2
-2
@@ -296,7 +296,7 @@ class Indexer(nn.Module):
296
cache_config=cache_config,
297
index_kpool=self.index_kpool,
298
)
299
- self.max_model_len = vllm_config.model_config.max_model_len
299
+ self.max_pool_len = vllm_config.model_config.max_model_len // self.index_kpool
300
self.prefix = prefix
301
from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
302
@@ -307,7 +307,7 @@ class Indexer(nn.Module):
307
self.scale_fmt,
308
self.topk_tokens,
309
self.head_dim,
310
- self.max_model_len,
310
+ self.max_pool_len,
311
self.max_total_seq_len,
312
self.topk_indices_buffer,
313
tail_cache=self.tail_cache,
vllm/models/glm5next/nvidia/sparse_indexer.py
+7
-7
@@ -117,7 +117,7 @@ def sparse_attn_indexer_kpool(
117
scale_fmt: str | None,
118
topk_tokens: int,
119
head_dim: int,
120
- max_model_len: int,
120
+ max_pool_len: int,
121
total_seq_lens: int,
122
topk_indices_buffer: torch.Tensor,
123
skip_k_cache_insert: bool,
@@ -153,7 +153,7 @@ def sparse_attn_indexer_kpool(
153
)
154
155
# Reserve profiler-visible memory for the worst-case decode logits,
156
- # whose shape is [B * next_n, max_model_len]. This profiling branch
156
+ # whose shape is [B * next_n, max_pool_len]. This profiling branch
157
# returns before invoking the logits kernel itself.
158
cfg = get_current_vllm_config_or_none()
159
worst_decode_tokens = 0
@@ -169,7 +169,7 @@ def sparse_attn_indexer_kpool(
169
sched.max_num_batched_tokens,
170
)
171
# float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
172
- decode_logits_elems = worst_decode_tokens * max_model_len * 4
172
+ decode_logits_elems = worst_decode_tokens * max_pool_len * 4
173
prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
174
max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
175
_ = torch.empty(
@@ -569,7 +569,7 @@ def sparse_attn_indexer_kpool(
569
seq_lens,
570
decode_metadata.block_table,
571
decode_metadata.schedule_metadata,
572
- max_model_len=max_model_len,
572
+ max_model_len=max_pool_len,
573
clean_logits=False,
574
)
575
num_rows = logits.shape[0]
@@ -668,7 +668,7 @@ class SparseAttnIndexerKpool(CustomOp):
668
scale_fmt: str,
669
topk_tokens: int,
670
head_dim: int,
671
- max_model_len: int,
671
+ max_pool_len: int,
672
max_total_seq_len: int,
673
topk_indices_buffer: torch.Tensor,
674
skip_k_cache_insert: bool = False,
@@ -682,7 +682,7 @@ class SparseAttnIndexerKpool(CustomOp):
682
self.scale_fmt = scale_fmt
683
self.topk_tokens = topk_tokens
684
self.head_dim = head_dim
685
- self.max_model_len = max_model_len
685
+ self.max_pool_len = max_pool_len
686
self.max_total_seq_len = max_total_seq_len
687
self.topk_indices_buffer = topk_indices_buffer
688
self.skip_k_cache_insert = skip_k_cache_insert
@@ -755,7 +755,7 @@ class SparseAttnIndexerKpool(CustomOp):
755
self.scale_fmt,
756
self.topk_tokens,
757
self.head_dim,
758
- self.max_model_len,
758
+ self.max_pool_len,
759
self.max_total_seq_len,
760
self.topk_indices_buffer,
761
self.skip_k_cache_insert,