perf(glm): bound prefill indexer logits workspace

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 分离部署
 ------------
 

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
 ----------------------------
 

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):

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()