diff --git a/vllm/models/glm5next/amd/sparse_indexer.py b/vllm/models/glm5next/amd/sparse_indexer.py index c9b68e973f..761f79d95d 100644 --- a/vllm/models/glm5next/amd/sparse_indexer.py +++ b/vllm/models/glm5next/amd/sparse_indexer.py @@ -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, diff --git a/vllm/models/glm5next/common/attention.py b/vllm/models/glm5next/common/attention.py index 21493174e8..4f776e8a8e 100644 --- a/vllm/models/glm5next/common/attention.py +++ b/vllm/models/glm5next/common/attention.py @@ -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, diff --git a/vllm/models/glm5next/nvidia/sparse_indexer.py b/vllm/models/glm5next/nvidia/sparse_indexer.py index 9dab659d4c..f7204949fc 100644 --- a/vllm/models/glm5next/nvidia/sparse_indexer.py +++ b/vllm/models/glm5next/nvidia/sparse_indexer.py @@ -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,