iHore/vLLM
8
0
Fork 0
Code Issues Pull Requests Actions 1 Packages Projects Releases Wiki Activity

[GLM5.3 Perf] Size the GLM-5 sparse indexer decode workspace, 3072 MiB GPU memory saved (#57701)

Browse Source 
Signed-off-by: yewentao256 <zhyanwentao@126.com>
This commit is contained in:
Wentao Ye
2026-09-19 16:57:01 -04:00
committed by GitHub
parent 133b71e0be
commit 36fa72d2d0
3 changed files with 17 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files
vllm/models/glm5next/amd/sparse_indexer.py
+8 -8
View File
@@ -100,7 +100,7 @@ def sparse_attn_indexer_kpool(
scale_fmt: str | None,
topk_tokens: int,
head_dim: int,
max_model_len: int,
max_pool_len: int,
total_seq_lens: int,
topk_indices_buffer: torch.Tensor,
skip_k_cache_insert: bool,
@@ -136,7 +136,7 @@ def sparse_attn_indexer_kpool(
)
# Reserve profiler-visible memory for the worst-case decode logits,
# whose shape is [B * next_n, max_model_len]. This profiling branch
# whose shape is [B * next_n, max_pool_len]. This profiling branch
# returns before invoking the logits kernel itself.
cfg = get_current_vllm_config_or_none()
worst_decode_tokens = 0
@@ -152,7 +152,7 @@ def sparse_attn_indexer_kpool(
sched.max_num_batched_tokens,
)
# float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
decode_logits_elems = worst_decode_tokens * max_model_len * 4
decode_logits_elems = worst_decode_tokens * max_pool_len * 4
prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
_ = torch.empty(
@@ -566,7 +566,7 @@ def sparse_attn_indexer_kpool(
seq_lens,
decode_metadata.block_table,
decode_metadata.schedule_metadata,
max_model_len=max_model_len,
max_model_len=max_pool_len,
)
num_rows = logits.shape[0]
# kpool: logits are pool-granular -> select topk_tokens//kpool pools,
@@ -645,7 +645,7 @@ class SparseAttnIndexerKpool(CustomOp):
scale_fmt: str,
topk_tokens: int,
head_dim: int,
max_model_len: int,
max_pool_len: int,
max_total_seq_len: int,
topk_indices_buffer: torch.Tensor,
skip_k_cache_insert: bool = False,
@@ -659,7 +659,7 @@ class SparseAttnIndexerKpool(CustomOp):
self.scale_fmt = scale_fmt
self.topk_tokens = topk_tokens
self.head_dim = head_dim
self.max_model_len = max_model_len
self.max_pool_len = max_pool_len
self.max_total_seq_len = max_total_seq_len
self.topk_indices_buffer = topk_indices_buffer
self.skip_k_cache_insert = skip_k_cache_insert
@@ -698,7 +698,7 @@ class SparseAttnIndexerKpool(CustomOp):
self.scale_fmt,
self.topk_tokens,
self.head_dim,
self.max_model_len,
self.max_pool_len,
self.max_total_seq_len,
self.topk_indices_buffer,
skip_k_cache_insert=self.skip_k_cache_insert,
@@ -715,7 +715,7 @@ class SparseAttnIndexerKpool(CustomOp):
self.scale_fmt,
self.topk_tokens,
self.head_dim,
self.max_model_len,
self.max_pool_len,
self.max_total_seq_len,
self.topk_indices_buffer,
self.skip_k_cache_insert,
vllm/models/glm5next/common/attention.py
+2 -2
View File
@@ -296,7 +296,7 @@ class Indexer(nn.Module):
cache_config=cache_config,
index_kpool=self.index_kpool,
)
self.max_model_len = vllm_config.model_config.max_model_len
self.max_pool_len = vllm_config.model_config.max_model_len // self.index_kpool
self.prefix = prefix
from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
@@ -307,7 +307,7 @@ class Indexer(nn.Module):
self.scale_fmt,
self.topk_tokens,
self.head_dim,
self.max_model_len,
self.max_pool_len,
self.max_total_seq_len,
self.topk_indices_buffer,
tail_cache=self.tail_cache,
vllm/models/glm5next/nvidia/sparse_indexer.py
+7 -7
View File
@@ -117,7 +117,7 @@ def sparse_attn_indexer_kpool(
scale_fmt: str | None,
topk_tokens: int,
head_dim: int,
max_model_len: int,
max_pool_len: int,
total_seq_lens: int,
topk_indices_buffer: torch.Tensor,
skip_k_cache_insert: bool,
@@ -153,7 +153,7 @@ def sparse_attn_indexer_kpool(
)
# Reserve profiler-visible memory for the worst-case decode logits,
# whose shape is [B * next_n, max_model_len]. This profiling branch
# whose shape is [B * next_n, max_pool_len]. This profiling branch
# returns before invoking the logits kernel itself.
cfg = get_current_vllm_config_or_none()
worst_decode_tokens = 0
@@ -169,7 +169,7 @@ def sparse_attn_indexer_kpool(
sched.max_num_batched_tokens,
)
# float32 logits -> 4 bytes/element; uint8 sentinel so elems == bytes.
decode_logits_elems = worst_decode_tokens * max_model_len * 4
decode_logits_elems = worst_decode_tokens * max_pool_len * 4
prefill_cap_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
max_logits_elems = max(decode_logits_elems, prefill_cap_elems)
_ = torch.empty(
@@ -569,7 +569,7 @@ def sparse_attn_indexer_kpool(
seq_lens,
decode_metadata.block_table,
decode_metadata.schedule_metadata,
max_model_len=max_model_len,
max_model_len=max_pool_len,
clean_logits=False,
)
num_rows = logits.shape[0]
@@ -668,7 +668,7 @@ class SparseAttnIndexerKpool(CustomOp):
scale_fmt: str,
topk_tokens: int,
head_dim: int,
max_model_len: int,
max_pool_len: int,
max_total_seq_len: int,
topk_indices_buffer: torch.Tensor,
skip_k_cache_insert: bool = False,
@@ -682,7 +682,7 @@ class SparseAttnIndexerKpool(CustomOp):
self.scale_fmt = scale_fmt
self.topk_tokens = topk_tokens
self.head_dim = head_dim
self.max_model_len = max_model_len
self.max_pool_len = max_pool_len
self.max_total_seq_len = max_total_seq_len
self.topk_indices_buffer = topk_indices_buffer
self.skip_k_cache_insert = skip_k_cache_insert
@@ -755,7 +755,7 @@ class SparseAttnIndexerKpool(CustomOp):
self.scale_fmt,
self.topk_tokens,
self.head_dim,
self.max_model_len,
self.max_pool_len,
self.max_total_seq_len,
self.topk_indices_buffer,
self.skip_k_cache_insert,