[None][perf] Speed up GLM-5.3-Flash serving on Blackwell

docs/source/deployment-guide/deployment-guide-for-glm-5.3-flash-on-trtllm.md

@@ -94,6 +94,10 @@ If you want to use the latest main branch, you can build from source: [https://n
 
 Use these configurations as starting points for 4x B200 with FP8 weights and a BF16 KV cache. Tune the batch size and token budget for your workload.
 
+For low-concurrency serving of long prompts, run the routed experts tensor-parallel with `--ep_size 1` instead of `--ep_size 4`. With prompts of more than 10K tokens, short outputs, MTP (`max_draft_len: 3`) and an FP8 KV cache, EP1 raised per-user output speed at concurrency 1 by about 6% on 4x B300, with concurrency-8 throughput about the same, because it avoids the expert-parallel load imbalance in the MoE all-reduce. This applies to the tensor-parallel configurations; keep `--ep_size 4` with attention data parallelism, which targets high concurrency, where expert-parallel load evens out. The chart in [Performance](#performance) shows where each configuration is fastest.
+
+For latency-sensitive serving, consider `TRTLLM_LOW_M_GEMM_BACKEND=auto` before launching the server. It routes BF16 projections with at most 32 rows to autotuned low-M GEMM kernels, so it is worth trying when decoding steps run with few tokens in flight. The kernel tuning it triggers adds minutes to every server start, so it suits long-running servers rather than ones that restart often.
+
 #### B200 FP8 Config
 
 ```bash
@@ -445,6 +449,6 @@ $$
 
 ## Performance
 
-The chart below shows the Pareto frontier of output-token throughput per GPU versus per-user output speed on 4x B200 with GLM-5.3-Flash FP8, a BF16 KV cache, and MTP enabled, at ISL 1024 / OSL 1024. The frontier is the efficient boundary of the measured TP4 / EP4 and attention DP4 / EP4 configurations with MTP: no other measured point is better on both axes. Other settings or workloads can give different results. Benchmarks were run using the `benchmark_serving` client above against TensorRT LLM 1.3.0rc29.
+The chart below shows the Pareto frontier of output-token throughput per GPU versus per-user output speed on 4x B300 with GLM-5.3-Flash FP8, an FP8 KV cache, and MTP enabled (`max_draft_len: 3`), at ISL 10240 / OSL 1024. The frontier is the efficient boundary of the measured TP4 / EP1, TP4 / EP4 and attention DP4 / EP4 configurations: no other measured point is better on both axes. TP4 / EP1 forms it up to concurrency 32 and attention DP4 / EP4 from concurrency 64. Other settings or workloads can give different results. Benchmarks were run with the aiperf client against a TensorRT LLM development build that includes these changes.
 
-![GLM-5.3-Flash FP8 Performance on 4x B200 with MTP](../media/glm_5_3_flash_fp8_perf.png)
+![GLM-5.3-Flash FP8 Performance on 4x B300 with MTP](../media/glm_5_3_flash_fp8_perf.png)

tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/backend.py

@@ -27,11 +27,14 @@
 
 from __future__ import annotations
 
+import os
+from collections.abc import Sequence
 from dataclasses import dataclass
 
 import torch
 
-from tensorrt_llm._utils import get_sm_version
+from tensorrt_llm._torch.pyexecutor.config_utils import unwrap_glm5_next_text_config
+from tensorrt_llm._utils import get_sm_version, is_sm_100f
 
 from ...interface import (
     AttentionForwardArgs,
@@ -46,6 +49,7 @@
     gather_fp8_kv_rows,
     kpool_expand,
     kpool_score,
+    kpool_score_fp8,
     kpool_update,
 )
 from .metadata import Glm5NextMamba2Metadata
@@ -136,6 +140,64 @@ class GlmKpoolSparseAttention(TrtllmAttention):
     # At the published 2112-entry selection this uses 132 MiB of BF16 rows.
     _FP8_QUERY_CHUNK_SIZE = 64
 
+    # FP32 elements one packed-context scoring call may hold: logits, their row
+    # copy and gathered pool keys grow with the visible pools, which a fresh-prompt
+    # profiling forward does not reach. plan_context_scores tiles calls to this
+    # bound and runtime_workspace_fixed_bytes reserves it in KV-cache sizing.
+    _CONTEXT_SCORE_ELEMENTS = 1 << 26
+
+    @staticmethod
+    def _context_score_cost(
+        rows: int, tables: int, pools: int, *, fp8: bool, head_dim: int, capacity_pools: int
+    ) -> int:
+        """FP32 elements a scoring call of ``rows`` rows over ``tables`` tables holds at once.
+
+        Counts the buffers that grow with the visible ``pools``: the BF16 scorer's
+        capacity-wide output, or every buffer :func:`kpool_score_fp8` keeps alive
+        together -- the padded logits, a multi-table call's gather temporary and
+        aligned row copy, and per gathered key its BF16 row, FP8 row, scale and
+        int64 slot index.
+        """
+        if not fp8 or pools == 0:
+            return rows * capacity_pools
+        # fp8_mqa_logits pads each row to a multiple of 256 columns, plus 256.
+        logits = -(-tables * pools // 256) * 256 + 256
+        copy = max(1024, -(-pools // 64) * 64) + pools if tables > 1 else 0
+        keys = tables * pools * (head_dim // 2 + head_dim // 4 + 3)
+        # int64 pool positions, pages and offsets.
+        return rows * (logits + copy) + keys + 6 * pools
+
+    @classmethod
+    def runtime_workspace_fixed_bytes(
+        cls,
+        model_config,
+        mapping,
+        *,
+        max_batch_size: int,
+        max_num_tokens: int,
+        max_seq_len: int,
+        enable_cuda_graph: bool,
+    ) -> int:
+        """Bytes of the largest scoring call plan_context_scores can emit for these capacities."""
+        if max_num_tokens <= 0 or max_seq_len <= 0:
+            return 0
+        config = unwrap_glm5_next_text_config(model_config.pretrained_config)
+        defaults = GlmKpoolSparseParams()
+        kpool = getattr(config, "index_kpool", defaults.index_kpool)
+        head_dim = getattr(config, "index_head_dim", defaults.index_head_dim)
+        pools = -(-max_seq_len // kpool)
+        tables = max(1, min(max_batch_size, max_num_tokens))
+
+        def cost(rows: int, tables: int, fp8: bool) -> int:
+            return cls._context_score_cost(
+                rows, tables, pools, fp8=fp8, head_dim=head_dim, capacity_pools=pools
+            )
+
+        largest = max(cost(max_num_tokens, tables, True), cost(max_num_tokens, 1, False))
+        # Calls fit the budget, except that a single row is never split.
+        bound = max(cls._CONTEXT_SCORE_ELEMENTS, cost(1, 1, True), cost(1, 1, False))
+        return min(largest, bound) * 4
+
     def __init__(
         self,
         layer_idx: int,
@@ -210,6 +272,9 @@ def __init__(
         #: construction). A float lets the kernel take it as an argument, without a
         #: per-call device multiply.
         self._native_bmm1_scale: float | None = None
+        #: Packed context rows score pools in FP8 with DeepGEMM MQA logits on SM100-family
+        #: GPUs; TRTLLM_GLM_INDEXER_BF16=1 keeps the BF16 Triton scorer for them.
+        self._indexer_fp8 = is_sm_100f() and os.environ.get("TRTLLM_GLM_INDEXER_BF16", "0") != "1"
 
     @classmethod
     def support_fused_rope(cls) -> bool:
@@ -329,19 +394,17 @@ def append_paged_state(
     def _rows(
         self,
         state: _GlmKpoolCacheState,
-        rows_per_request: int = 1,
         request_ids: torch.Tensor | None = None,
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
-        """Select generation tables, or all tables for packed context request IDs.
+        """Select generation tables, or all tables for per-row request IDs.
 
-        Verification repeats generation tables per token; callers supply each
-        query's visible length for verification and packed context rows.
+        Rows with request_ids (packed context, or verification rows given their
+        tables' batch rows) index all tables and supply each query's visible
+        length; single-token decode rows use the generation tables and live lengths.
         """
         if request_ids is not None:
             return state.block_tables, None
         gen = slice(state.num_contexts, None)
-        if rows_per_request > 1:
-            return state.block_tables[gen].repeat_interleave(rows_per_request, dim=0), None
         return state.block_tables[gen], state.kv_lens[gen]
 
     def update_pool_keys(
@@ -372,6 +435,63 @@ def update_pool_keys(
             request_ids=request_ids,
         )
 
+    def plan_context_scores(
+        self, metadata, spans: Sequence[tuple[int, int]]
+    ) -> list[tuple[int, int, int, int, int]]:
+        """Split packed context scoring into calls within ``_CONTEXT_SCORE_ELEMENTS``.
+
+        spans lists ``(rows, visible bound)`` of each context request in packed
+        order. Each returned ``(row_start, row_end, table_start, table_end,
+        max_kv_len)`` scores rows ``row_start:row_end`` against context tables
+        ``table_start:table_end``. Consecutive requests share a call while it
+        fits; a request that does not fit alone is split by rows. Costs count
+        only the transients that grow with the visible pools.
+        """
+        state = self._cache_state(metadata)
+        kpool = self.sparse_params.index_kpool
+        budget = self._CONTEXT_SCORE_ELEMENTS
+        # The BF16 scorer writes every pool column of the configured capacity.
+        capacity_pools = -(-state.block_tables.shape[1] * state.tokens_per_block // kpool)
+
+        def cost(rows: int, tables: int, pools: int) -> int:
+            return self._context_score_cost(
+                rows,
+                tables,
+                pools,
+                fp8=self._indexer_fp8,
+                head_dim=self.sparse_params.index_head_dim,
+                capacity_pools=capacity_pools,
+            )
+
+        calls: list[tuple[int, int, int, int, int]] = []
+        group: list[int] | None = None  # [row_start, rows, table_start, tables, max visible]
+        row_start = 0
+        for table, (rows, visible) in enumerate(spans):
+            if group is not None:
+                merged = max(group[4], visible)
+                if cost(group[1] + rows, group[3] + 1, merged // kpool) <= budget:
+                    group[1] += rows
+                    group[3] += 1
+                    group[4] = merged
+                    row_start += rows
+                    continue
+                calls.append(
+                    (group[0], group[0] + group[1], group[2], group[2] + group[3], group[4])
+                )
+                group = None
+            if cost(rows, 1, visible // kpool) <= budget:
+                group = [row_start, rows, table, 1, visible]
+            else:
+                fixed = cost(0, 1, visible // kpool)
+                tile = max(1, (budget - fixed) // max(1, cost(1, 1, visible // kpool) - fixed))
+                for start in range(row_start, row_start + rows, tile):
+                    end = min(start + tile, row_start + rows)
+                    calls.append((start, end, table, table + 1, visible))
+            row_start += rows
+        if group is not None:
+            calls.append((group[0], group[0] + group[1], group[2], group[2] + group[3], group[4]))
+        return calls
+
     def score_pools(
         self,
         q: torch.Tensor,
@@ -384,21 +504,47 @@ def score_pools(
         rows_per_request: int = 1,
         request_ids: torch.Tensor | None = None,
         max_kv_len: int | None = None,
+        table_range: tuple[int, int] | None = None,
     ) -> torch.Tensor:
         """Score complete pools into FP32 [rows, pool_capacity].
 
-        q is [rows, heads, head_dim], weights is [rows, heads]. Packed context uses
-        request_ids; verification repeats generation tables rows_per_request times.
-        Both supply per-query kv_lens. Single-token decode uses live metadata lengths.
+        q is [rows, heads, head_dim], weights is [rows, heads]. Packed context and
+        verification (rows_per_request rows per request) use request_ids and
+        per-query kv_lens. Single-token decode uses live metadata lengths.
         Incomplete/invisible pools receive the FP32 minimum before top-k. Eager
         callers that know a host bound on the visible lengths pass max_kv_len so
         only reachable pools are scored; columns past it are left unwritten and
-        the top-k must scan only each row's complete pools.
+        the top-k must scan only each row's complete pools. Packed context rows
+        with such a bound use the FP8 scorer on SM100-family GPUs, whose result
+        may be a strided view. table_range, from :meth:`plan_context_scores`,
+        names the context tables those rows' request_ids fall in.
         """
+        if rows_per_request > 1 and request_ids is None:
+            raise ValueError("glm_kpool verification rows need request_ids")
         state = self._cache_state(metadata)
-        tables, gen_lens = self._rows(state, rows_per_request, request_ids=request_ids)
+        tables, gen_lens = self._rows(state, request_ids)
         capacity = state.block_tables.shape[1] * state.tokens_per_block
         kpool = self.sparse_params.index_kpool
+        packed_context = request_ids is not None and rows_per_request == 1
+        if self._indexer_fp8 and packed_context and max_kv_len is not None:
+            # Packed context rows only reference the context requests' tables.
+            first, last = (0, state.num_contexts) if table_range is None else table_range
+            scores = kpool_score_fp8(
+                q,
+                weights,
+                state.index_pool,
+                tables[first:last],
+                kv_lens,
+                state.tokens_per_block,
+                head_dim=self.sparse_params.index_head_dim,
+                kpool=kpool,
+                q_scale=q_scale,
+                w_scale=w_scale,
+                request_ids=request_ids - first if first else request_ids,
+                max_kv_len=max_kv_len,
+            )
+            if scores is not None:
+                return scores
         return kpool_score(
             q,
             weights,
@@ -415,9 +561,10 @@ def score_pools(
             # dot is an fp32 accumulation of exact products, like tf32 on the
             # widened operands (measured: identical selections).
             precision="bf16",
-            # Context rows: the query tokens of a request share its block
-            # table, so a program gathers each pool-key block once for 16 rows.
-            rows_per_program=1 if request_ids is None else 16,
+            # Packed context rows share their request's block table, so a program
+            # gathers each pool-key block once for 16 rows. Verification batches
+            # are small; one row per program keeps more programs in flight.
+            rows_per_program=16 if request_ids is not None and rows_per_request == 1 else 1,
             request_ids=request_ids,
             max_kv_len=max_kv_len,
         )
@@ -428,7 +575,6 @@ def expand_selection(
         metadata,
         *,
         kv_lens: torch.Tensor | None = None,
-        rows_per_request: int = 1,
         request_ids: torch.Tensor | None = None,
     ) -> torch.Tensor:
         """Selected pools ``[N, select_k]`` -> latent row ids for the kernel.
@@ -439,7 +585,7 @@ def expand_selection(
         row addressing as :meth:`score_pools`.
         """
         state = self._cache_state(metadata)
-        tables, gen_lens = self._rows(state, rows_per_request, request_ids=request_ids)
+        tables, gen_lens = self._rows(state, request_ids)
         _, base_row, rows_per_slot = latent_pool_rows(state.latent_pool)
         return kpool_expand(
             selected,

tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/kernels.py

@@ -512,6 +512,71 @@ def kpool_score(
     return out
 
 
+def kpool_score_fp8(
+    q: torch.Tensor,
+    weights: torch.Tensor,
+    index_pool: torch.Tensor,
+    block_tables: torch.Tensor,
+    kv_lens: torch.Tensor,
+    tokens_per_block: int,
+    *,
+    head_dim: int,
+    kpool: int,
+    q_scale: float,
+    w_scale: float,
+    request_ids: torch.Tensor,
+    max_kv_len: int,
+) -> torch.Tensor | None:
+    """FP8 :func:`kpool_score` for packed context rows, via DeepGEMM ``fp8_mqa_logits``.
+
+    The pool keys of every block table are gathered into ``[tables * pools, head_dim]``
+    and quantized per row; queries are quantized per (row, head), with their scales
+    folded into the head weights, so each score is ``sum_h w * relu(q . key)`` as in
+    :func:`kpool_score`. ``max_kv_len`` bounds every row's ``kv_lens``. Row ``r``
+    scores only the first ``kv_lens[r] // kpool`` pools of table ``request_ids[r]``;
+    later columns are not meaningful, so readers must stop there. The result may be
+    a strided view with an aligned row stride. Returns ``None`` when there are no
+    rows or no complete pools.
+    """
+    from tensorrt_llm.deep_gemm import fp8_mqa_logits
+
+    n = q.shape[0]
+    pools = max(0, int(max_kv_len)) // kpool
+    tables = block_tables.shape[0]
+    if n == 0 or pools == 0:
+        return None
+    positions = torch.arange(pools, device=q.device) * kpool
+    page = positions // tokens_per_block
+    offset = positions - page * tokens_per_block
+    slots = block_tables[:, page].long()  # [tables, pools]
+    keys = index_pool[slots, offset, 2 * head_dim : 3 * head_dim].reshape(-1, head_dim)
+    half = head_dim // 2
+    keys_fp8, key_scale = torch.ops.trtllm.fused_cat_fp8(keys[:, :half], keys[:, half:], False)
+    q_rows = q.reshape(-1, head_dim)
+    q_fp8, q_row_scale = torch.ops.trtllm.fused_cat_fp8(q_rows[:, :half], q_rows[:, half:], False)
+    head_weights = weights.float() * q_row_scale.view(n, -1) * (w_scale * q_scale)
+    start = request_ids.to(torch.int32) * pools
+    end = start + (kv_lens // kpool).to(torch.int32)
+    logits = fp8_mqa_logits(
+        q_fp8.view(q.shape),
+        (keys_fp8, key_scale.view(-1)),
+        head_weights,
+        start,
+        end,
+        clean_logits=False,
+    )
+    if tables == 1:
+        # DeepGEMM pads the logits rows to an aligned stride, as the pool top-k needs.
+        return logits[:, :pools]
+    # Keep the copy's rows aligned and wider than the selection, like the logits arena.
+    out = torch.empty(
+        n, max(1024, triton.cdiv(pools, 64) * 64), dtype=torch.float32, device=q.device
+    )
+    rows = torch.arange(n, device=q.device)
+    out[:, :pools] = logits[:, : tables * pools].view(n, tables, pools)[rows, request_ids.long()]
+    return out
+
+
 @triton.jit
 def _kpool_expand_kernel(
     SEL,

tensorrt_llm/_torch/attention/backends/sparse/glm_kpool/native_decode.py

@@ -29,6 +29,8 @@ def __init__(self) -> None:
         # Largest query rows and heads the sizes cover.
         self._covered_rows = 0
         self._covered_heads = 0
+        # Device zero for discarding empty rows; a Python 0 would be filled per call.
+        self._zero: torch.Tensor | None = None
 
     def prepare_workspace(
         self,
@@ -70,6 +72,10 @@ def prepare_workspace(
                 )
             # TRTLLM-GEN resets its counters after each invocation.
             self._counter = torch.zeros(self._counter_size, dtype=torch.uint8, device=q.device)
+        if self._zero is None:
+            if capturing:
+                raise RuntimeError("glm_kpool native decode zero must be allocated before capture")
+            self._zero = torch.zeros((), dtype=torch.bfloat16, device=q.device)
 
     def __call__(
         self,
@@ -112,4 +118,4 @@ def __call__(
         ).squeeze(1)
         # Empty (padded) requests use a dummy selection to keep the kernel's
         # length positive; discard that output, matching FlashMLA's empty row.
-        return torch.where(nonempty[:, None, None], output, 0)
+        return torch.where(nonempty[:, None, None], output, self._zero)