[None][feat] Add a glm_kpool sparse attention config

tensorrt_llm/llmapi/llm_args.py

@@ -1724,6 +1724,35 @@ def to_sparse_metadata_params(self, **kwargs):
             layout=CSA2Layout.from_hf_config(kwargs.get("pretrained_config")))
 
 
+class GlmKpoolSparseAttentionConfig(BaseSparseAttentionConfig):
+    """Use checkpoint-owned GLM-5.3-Flash k-pool indexer geometry with the PyTorch sparse-MLA backend.
+
+    The indexer top-k, pool size and key width come from the checkpoint.
+    """
+    algorithm: Literal["glm_kpool"] = "glm_kpool"
+
+    def supports_backend(self, backend: str) -> bool:
+        return backend == "pytorch"
+
+    def to_sparse_params(self, **kwargs):
+        from tensorrt_llm._torch.attention.backends.sparse.glm_kpool.params import \
+            GlmKpoolSparseParams
+        from tensorrt_llm._torch.pyexecutor.config_utils import \
+            unwrap_glm5_next_text_config
+
+        pretrained_config = kwargs.get("pretrained_config")
+        if pretrained_config is None:
+            return GlmKpoolSparseParams()
+        config = unwrap_glm5_next_text_config(pretrained_config)
+        return GlmKpoolSparseParams(
+            index_topk=int(config.index_topk),
+            index_kpool=int(config.index_kpool),
+            index_always_select_tail=bool(
+                config.index_kpool_always_select_tail),
+            index_head_dim=int(config.index_head_dim),
+        )
+
+
 class SkipSoftmaxAttentionConfig(BaseSparseAttentionConfig):
     """Configuration for skip softmax attention."""
     algorithm: Literal["skip_softmax"] = Field(default="skip_softmax")
@@ -4421,6 +4450,7 @@ def _to_pybind(self):
         DeepSeekSparseAttentionConfig,
         DeepSeekV4SparseAttentionConfig,
         CSA2SparseAttentionConfig,
+        GlmKpoolSparseAttentionConfig,
         SkipSoftmaxAttentionConfig,
         MiniMaxM3SparseAttentionConfig,
     ],

tensorrt_llm/usage/llm_args_golden_manifest.json

@@ -1380,6 +1380,7 @@
         "csa2",
         "dsa",
         "deepseek_v4",
+        "glm_kpool",
         "minimax_m3",
         "qsa",
         "rocket",

tests/unittest/_torch/modeling/test_glm5_next_contracts.py

@@ -752,6 +752,31 @@ def test_cache_manager_routing_guards(route, monkeypatch):
             get_kv_cache_manager_cls(config, kv, True, transceiver)
 
 
+@pytest.mark.cpu_only
+def test_glm_kpool_sparse_config_lowers_the_checkpoint_geometry():
+    from tensorrt_llm._torch.attention.backends.sparse.glm_kpool import (
+        GlmKpoolSparseAttention,
+        GlmKpoolSparseParams,
+    )
+    from tensorrt_llm._torch.attention.backends.utils import get_attention_backend
+    from tensorrt_llm._torch.pyexecutor.config_utils import unwrap_glm5_next_text_config
+    from tensorrt_llm.llmapi.llm_args import GlmKpoolSparseAttentionConfig
+
+    pretrained = _config()
+    text = unwrap_glm5_next_text_config(pretrained)
+    sparse = GlmKpoolSparseAttentionConfig()
+    params = sparse.to_sparse_params(pretrained_config=pretrained)
+    assert params == GlmKpoolSparseParams(
+        index_topk=text.index_topk,
+        index_kpool=text.index_kpool,
+        index_always_select_tail=text.index_kpool_always_select_tail,
+        index_head_dim=text.index_head_dim,
+    )
+    assert sparse.to_sparse_params(pretrained_config=text) == params
+    assert sparse.supports_backend("pytorch") and not sparse.supports_backend("tensorrt")
+    assert get_attention_backend("TRTLLM", sparse_params=params) is GlmKpoolSparseAttention
+
+
 @pytest.mark.cpu_only
 @pytest.mark.parametrize("fp8_kv_cache", [False, True], ids=["bf16-kv", "fp8-kv"])
 @pytest.mark.parametrize(

tests/unittest/api_stability/references/llm.yaml

@@ -323,7 +323,7 @@ methods:
         default: null
         status: prototype
       sparse_attention_config:
-        annotation: Union[tensorrt_llm.llmapi.llm_args.QSASparseAttentionConfig, tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekV4SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.CSA2SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig, NoneType]
+        annotation: Union[tensorrt_llm.llmapi.llm_args.QSASparseAttentionConfig, tensorrt_llm.llmapi.llm_args.RocketSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.DeepSeekV4SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.CSA2SparseAttentionConfig, tensorrt_llm.llmapi.llm_args.GlmKpoolSparseAttentionConfig, tensorrt_llm.llmapi.llm_args.SkipSoftmaxAttentionConfig, tensorrt_llm.llmapi.llm_args.MiniMaxM3SparseAttentionConfig, NoneType]
         default: null
         status: prototype
       kv_cache_compression_config:

tests/unittest/usage/test_llmapi_config_capture.py

@@ -741,6 +741,7 @@ def test_collect_llm_api_config_captures_sparse_algorithm_for_every_arm():
         CSA2SparseAttentionConfig,
         DeepSeekSparseAttentionConfig,
         DeepSeekV4SparseAttentionConfig,
+        GlmKpoolSparseAttentionConfig,
         MiniMaxM3SparseAttentionConfig,
         QSASparseAttentionConfig,
         RocketSparseAttentionConfig,
@@ -753,6 +754,7 @@ def test_collect_llm_api_config_captures_sparse_algorithm_for_every_arm():
         DeepSeekSparseAttentionConfig(),
         DeepSeekV4SparseAttentionConfig(),
         CSA2SparseAttentionConfig(),
+        GlmKpoolSparseAttentionConfig(),
         RocketSparseAttentionConfig(),
         SkipSoftmaxAttentionConfig(),
         MiniMaxM3SparseAttentionConfig(),