From 6388f7b34b3da603664cafacc03858523ab2f51c Mon Sep 17 00:00:00 2001 From: shihaobai <42648726+shihaobai@users.noreply.github.com> Date: Tue, 3 Nov 2026 14:49:48 +0800 Subject: [PATCH] perf(glm): bound prefill indexer logits workspace --- .../source/tutorial/fp8_kv_quantization.rst | 5 + .../source/tutorial/fp8_kv_quantization.rst | 6 ++ lightllm/models/glm5_next/indexer.py | 22 ++++- .../glm5_next/test_indexer_workspace.py | 94 +++++++++++++++++++ 4 files changed, 123 insertions(+), 4 deletions(-) create mode 100644 unit_tests/models/glm5_next/test_indexer_workspace.py diff --git a/docs/CN/source/tutorial/fp8_kv_quantization.rst b/docs/CN/source/tutorial/fp8_kv_quantization.rst index 868598caed..8b4817d06d 100644 --- a/docs/CN/source/tutorial/fp8_kv_quantization.rst +++ b/docs/CN/source/tutorial/fp8_kv_quantization.rst @@ -66,6 +66,11 @@ DeepSeek-V3.2 和 GLM-5.3 Flash 使用 ``--llm_kv_type fp8kv_dsa``，无需提 --model_dir /path/to/GLM-5.3-Flash --tp 4 \ --data_type bfloat16 --llm_kv_type fp8kv_dsa +GLM-5.3 Flash 的 indexer 默认将每块 FP32 logits 临时张量限制在 64 MiB。 +可通过 ``LIGHTLLM_GLM53_PREFILL_LOGITS_MIB`` 设置 16–1024 MiB 的整数预算； +更大的预算允许每块并行处理更多 query，同时增加临时显存占用。 +该设置也适用于未量化 KV，只控制 indexer logits 分块，不限制模型的总显存。 + P/D 分离部署 ------------ diff --git a/docs/EN/source/tutorial/fp8_kv_quantization.rst b/docs/EN/source/tutorial/fp8_kv_quantization.rst index 33c9ebfc8f..728893bb3d 100644 --- a/docs/EN/source/tutorial/fp8_kv_quantization.rst +++ b/docs/EN/source/tutorial/fp8_kv_quantization.rst @@ -70,6 +70,12 @@ The following example uses GLM-5.3 Flash, which requires ``bfloat16`` and suppor --model_dir /path/to/GLM-5.3-Flash --tp 4 \ --data_type bfloat16 --llm_kv_type fp8kv_dsa +The GLM-5.3 Flash indexer limits each temporary FP32 logits chunk to 64 MiB by default. +Set ``LIGHTLLM_GLM53_PREFILL_LOGITS_MIB`` to an integer from 16 to 1024 MiB to change this +budget. Larger budgets allow more queries per chunk at the cost of more temporary GPU +memory. This also applies to unquantized KV; it controls indexer logits chunking, not +the model's total GPU memory use. + Disaggregated P/D Deployment ---------------------------- diff --git a/lightllm/models/glm5_next/indexer.py b/lightllm/models/glm5_next/indexer.py index 5df510e5de..0100539cc9 100644 --- a/lightllm/models/glm5_next/indexer.py +++ b/lightllm/models/glm5_next/indexer.py @@ -1,13 +1,24 @@ -from lightllm.common.attention_state.full_att import FullAttentionState +import os + import torch import triton +from lightllm.common.attention_state.full_att import FullAttentionState from lightllm.utils.vllm_utils import HAS_VLLM, vllm_ops from .triton_kernel.index_quant import hadamard_transform_quant_fp8 from .triton_kernel.kpool import compress_pools, gather_pools, gather_paged_pools, get_pool_ranges, expand_topk +def get_prefill_logits_chunk_size(num_q_tokens: int, max_pools: int, logits_budget_mib: int) -> int: + """Bound the FP32 temporary without changing the per-query indexer math.""" + if not 16 <= logits_budget_mib <= 1024: + raise ValueError("GLM prefill logits budget must be between 16 and 1024 MiB") + if max_pools <= 0: + raise ValueError("max_pools must be positive") + return max(1, min(num_q_tokens, logits_budget_mib * 1024 * 1024 // 4 // max_pools)) + + class Glm5NextNsaInfer: """K-pool indexing with pooled token KV and a small per-request raw tail.""" @@ -18,6 +29,8 @@ def __init__(self, layer_idx, network_config, tp_world_size): self.dim = network_config["index_head_dim"] self.eps = network_config["rms_norm_eps"] self.pool_size = network_config["index_kpool"] + self.prefill_logits_budget_mib = int(os.environ.get("LIGHTLLM_GLM53_PREFILL_LOGITS_MIB", "64")) + get_prefill_logits_chunk_size(1, 1, self.prefill_logits_budget_mib) def select_topk_indices(self, logits, lengths, indices): """Select row-relative indices, with valid entries before -1 padding.""" @@ -96,9 +109,8 @@ def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_s att_state.lengths, infer_state.b1_cu_q_seq_len, infer_state.max_q_seq_len, max_pools ) groups = torch.empty((q_fp8.shape[0], self.topk // 4), dtype=torch.int32, device=q_fp8.device) - # Budget 64 MiB for FP32 logits. At 1M tokens this permits 64 queries per chunk, - # limiting Q parallelism even when the prefill batch contains many tokens. - chunk_size = max(1, min(q_fp8.shape[0], 16 * 1024 * 1024 // max_pools)) + # Keep the 64 MiB default; larger budgets trade memory for query parallelism. + chunk_size = get_prefill_logits_chunk_size(q_fp8.shape[0], max_pools, self.prefill_logits_budget_mib) import deep_gemm for start in range(0, q_fp8.shape[0], chunk_size): @@ -113,6 +125,8 @@ def _get_prefill_indices(self, q_fp8, weights, packed_buffer, infer_state, att_s max_seqlen_k=max_pools, ) self.select_topk_indices(logits, lengths[start:end], groups[start:end]) + # Release this chunk before DeepGEMM allocates the next one. + del logits return groups def _get_decode_indices(self, q_fp8, weights, packed_buffer, infer_state, att_state, max_pools): diff --git a/unit_tests/models/glm5_next/test_indexer_workspace.py b/unit_tests/models/glm5_next/test_indexer_workspace.py new file mode 100644 index 0000000000..22cc1d9450 --- /dev/null +++ b/unit_tests/models/glm5_next/test_indexer_workspace.py @@ -0,0 +1,94 @@ +import weakref +from types import SimpleNamespace + +import pytest +import torch + +from lightllm.models.glm5_next.indexer import Glm5NextNsaInfer, get_prefill_logits_chunk_size + + +@pytest.mark.parametrize("budget_mib,expected", [(64, 64), (128, 128), (256, 256), (512, 512)]) +def test_one_million_context_query_parallelism(budget_mib, expected): + assert get_prefill_logits_chunk_size(2048, 262144, budget_mib) == expected + + +@pytest.mark.parametrize("num_q_tokens,max_pools", [(1, 262144), (8, 8192), (2048, 640), (8192, 262144), (0, 262144)]) +def test_default_matches_previous_bound(num_q_tokens, max_pools): + assert get_prefill_logits_chunk_size(num_q_tokens, max_pools, 64) == max( + 1, min(num_q_tokens, 16 * 1024 * 1024 // max_pools) + ) + + +@pytest.mark.parametrize("budget_mib", [0, -1, 15, 1025]) +def test_reject_unbounded_workspace(budget_mib): + with pytest.raises(ValueError, match="budget"): + get_prefill_logits_chunk_size(2048, 262144, budget_mib) + + +def test_environment_controls_each_indexer(monkeypatch): + monkeypatch.setenv("LIGHTLLM_GLM53_PREFILL_LOGITS_MIB", "256") + indexer = Glm5NextNsaInfer( + 0, {"index_kpool": 4, "index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5}, 1 + ) + assert indexer.prefill_logits_budget_mib == 256 + + +@pytest.mark.parametrize("max_pools", [0, -1]) +def test_reject_invalid_pool_capacity(max_pools): + with pytest.raises(ValueError, match="max_pools"): + get_prefill_logits_chunk_size(128, max_pools, 64) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="requires the GLM GPU indexer kernels") +def test_workspace_preserves_mixed_request_pool_selection(monkeypatch): + """Exercise an unaligned request boundary and a full 1M key range.""" + if torch.cuda.mem_get_info()[0] < 1536 * 1024 ** 2: + pytest.skip("requires at least 1.5 GiB of unused GPU memory") + import deep_gemm + + fp8_mqa_logits = deep_gemm.fp8_mqa_logits + previous_logits = None + + def track_logits_lifetime(*args, **kwargs): + nonlocal previous_logits + if previous_logits is not None: + assert previous_logits() is None, "the previous logits chunk is still alive" + logits = fp8_mqa_logits(*args, **kwargs) + previous_logits = weakref.ref(logits) + return logits + + monkeypatch.setattr(deep_gemm, "fp8_mqa_logits", track_logits_lifetime) + monkeypatch.setenv("LIGHTLLM_GLM53_PREFILL_LOGITS_MIB", "64") + torch.manual_seed(20261029) + max_pools = 262144 + packed_buffer = torch.empty((max_pools * 2, 1, 132), dtype=torch.uint8, device="cuda") + packed_buffer[:, 0, :128] = torch.randn(max_pools * 2, 128, device="cuda").to(torch.float8_e4m3fn).view(torch.uint8) + packed_buffer[:, 0, 128:] = torch.ones(max_pools * 2, device="cuda").view(torch.uint8).view(-1, 4) + req_to_token_indexs = torch.full((2, 1048576), -1, dtype=torch.int32, device="cuda") + for batch in range(2): + req_to_token_indexs[batch, 3::4] = torch.arange( + batch * max_pools, (batch + 1) * max_pools, dtype=torch.int32, device="cuda" + ) + q_fp8 = torch.randn(128, 32, 128, device="cuda").to(torch.float8_e4m3fn) + weights = torch.randn(128, 32, device="cuda") * 0.05 + lengths = torch.cat( + [torch.arange(65536 - 30, 65537, device="cuda"), torch.arange(1048576 - 96, 1048577, device="cuda")] + ).to(torch.int32) + infer_state = SimpleNamespace( + b_req_idx=torch.tensor([0, 1], dtype=torch.int32, device="cuda"), + b_seq_len=torch.tensor([65536, 1048576], dtype=torch.int32, device="cuda"), + b1_cu_q_seq_len=torch.tensor([0, 31, 128], dtype=torch.int32, device="cuda"), + max_q_seq_len=97, + req_manager=SimpleNamespace(req_to_token_indexs=req_to_token_indexs), + ) + att_state = SimpleNamespace(lengths=lengths) + indexer = Glm5NextNsaInfer( + 0, {"index_kpool": 4, "index_topk": 2048, "index_n_heads": 32, "index_head_dim": 128, "rms_norm_eps": 1e-5}, 1 + ) + expected = indexer._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + expected = expected.sort(dim=1).values + for budget in [256, 512]: + indexer.prefill_logits_budget_mib = budget + actual = indexer._get_prefill_indices(q_fp8, weights, packed_buffer, infer_state, att_state, max_pools) + torch.testing.assert_close(actual.sort(dim=1).values, expected, rtol=0, atol=0) + assert ((actual >= 0) & (actual < (lengths // 4)[:, None])).all().item()