Repository navigation
[None][feat] Add a glm_kpool sparse attention config - #20811
Closed
ruocheng-nv wants to merge 1 commit into
Closed
ruocheng-nv wants to merge 1 commit into
ruocheng-nv wants to merge 1 commit into
Conversation
GLM-5.3-Flash's sparse-MLA layers build their glm_kpool backend inside the model, so there is no sparse_attention_config that names it. Add GlmKpoolSparseAttentionConfig (algorithm "glm_kpool") to the sparse attention configs. Like CSA2, it has no geometry fields: to_sparse_params reads the indexer top-k, pool size, tail selection and key width from the GLM-5.3-Flash text config and lowers to GlmKpoolSparseParams, which the existing registry dispatches to GlmKpoolSparseAttention. Nothing fills the config yet, so serving behavior is unchanged. Signed-off-by: Ruocheng Jia <ruochengj@nvidia.com>
ruocheng-nv
force-pushed
the
feat/glm53-kpool-sparse-config
branch
from
October 22, 2026 08:30
8f91ca3 to
4587521
Compare
ruocheng-nv
changed the title
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
GLM-5.3-Flash's sparse-MLA layers build their
glm_kpoolbackend inside the model (create_attention(..., sparse_params=GlmKpoolSparseParams(...))), so nosparse_attention_confignames it and the engine resolves plainTrtllmAttention.This PR only adds the config:
GlmKpoolSparseAttentionConfiginllm_args.py, added to theSparseAttentionConfigunion.algorithm: "glm_kpool"is its only field. LikeCSA2SparseAttentionConfig, the geometry comes from the checkpoint:to_sparse_params(pretrained_config=...)readsindex_topk,index_kpool,index_kpool_always_select_tailandindex_head_dimfrom the GLM-5.3-Flash text config (composite or flattened) and returnsGlmKpoolSparseParams. The existing sparse registry already dispatches those params toGlmKpoolSparseAttention.supports_backendis PyTorch only.Nothing fills the config yet, so serving behavior is unchanged. A follow-up can have
ModelConfigfill it for GLM-5.3-Flash so the engine dispatches to the same backend the layers use; that needs the hybrid-model sparse allow-list to acceptglm_kpooland the vision tower to drop the config.Generated files and approvals
tensorrt_llm/usage/llm_args_golden_manifest.json:glm_kpooladded tosparse_attention_config.algorithm's allowed values (scripts/generate_llm_args_golden_manifest.py). Needs telemetry/privacy CODEOWNER approval.tensorrt_llm/schemas/trtllm-serve-config.schema.json,trtllm-serve-disagg-config.schema.json: regenerated (scripts/generate_trtllm_serve_schemas.py).tests/unittest/api_stability/references/llm.yaml: union annotation updated.Test Coverage
tests/unittest/_torch/modeling/test_glm5_next_contracts.py::test_glm_kpool_sparse_config_lowers_the_checkpoint_geometry: lowering from the composite and the flattened config, backend support, and registry dispatch toGlmKpoolSparseAttention.tests/unittest/usage/test_llmapi_config_capture.py: the new union arm is captured.test_glm5_next_contracts.py,usage/test_llmapi_config_capture.py,api_stability/,llmapi/test_trtllm_serve_config_schemas.py);generate_trtllm_serve_schemas.py --checkis clean. These ran on a slightly older main with the identical change, because the local build predates the current main bindings.PR Checklist