[Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU by jayhenry · Pull Request #2108 · InternLM/xtuner

jayhenry · 2026-09-23T06:54:44Z

Stack (bottom to top):

[Feature] Add GLM-5.3-Flash F0: 25B cropped reference checkpoint builder #2105 feat/glm53flash-materialize-full-f0 → main
[Feature] Add GLM-5.3-Flash F3: Kimi Delta Attention (KDA) #2106 feat/glm53flash-f3-kda → feat/glm53flash-materialize-full-f0
[Feature] Add GLM-5.3-Flash F4: mHC four-stream residual #2107 feat/glm53flash-f4-mhc → feat/glm53flash-f3-kda
[Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU #2108 feat/glm53flash-f5-nope-dsa → feat/glm53flash-f4-mhc ← you are here
[Feature] Add GLM-5.3-Flash F1: VL data preprocessing pipeline #2109 feat/glm53flash-f1-vl-data → feat/glm53flash-f5-nope-dsa
[Feature] Add GLM-5.3-Flash F2: vision tower + projector (eager) #2110 feat/glm53flash-f2-vision-tower → feat/glm53flash-f1-vl-data
[Feature] Add GLM-5.3-Flash F6 core: text model + MTP + compose model #2111 feat/glm53flash-f6-text-moe → feat/glm53flash-f2-vision-tower

Base is #2107's branch (layer 3). Review only this PR's own diff.

Summary

Stack layer 4/7 of GLM-5.3-Flash support (base: layer 3, F4 mHC).

Implements the F5 milestone: the NoPE (qk_rope_head_dim=0) DeepSeek Sparse Attention layers, their KPool indexer (pools of index_kpool consecutive tokens scored together instead of per-token top-k), a new flash_mla_cudnn SparseMLA backend (FlashMLA forward + existing cuDNN backward), and the clamped SwiGLU activation used by GLM-5.3-Flash's dense/shared/MoE MLPs.

Key pieces:

xtuner/v1/ops/act_fn.py: native_clamped_swiglu + MoEActFnConfig support, wired through DenseMLP/MoEMLP.
xtuner/v1/ops/sparse_mla/kpool.py: pool layout, causal ranges, and top-k pool selection (torch reference + TileLang-backed production path).
xtuner/v1/ops/sparse_mla/flash_mla_cudnn.py: new SparseMLA backend.
xtuner/v1/model/moe/glm53/nope_dsa_mla.py: NoPEDSAMLAConfig, KPoolIndexer, NoPEDSAMultiLatentAttention.
xtuner/v1/ops/sparse_mla/protocol.py: KPoolIndexerBackend/KPoolTopKIndicesProtocol (GLM-5.3-Flash's KPool only supports torch/tilelang, unlike GLM-5.2's 6-way DSAIndexerBackend).

Also folds in what were originally two separate follow-up commits, now part of this PR:

DSA backends get explicit defaults instead of inheriting: indexer_backend used to default to None and fall back to sparse_mla_backend, but the two name different, only partially-overlapping backend vocabularies (flash_mla_cudnn is SparseMLA-only; deep_gemm_fp8/cute_dsl are indexer-only), so the fallback could hand the indexer factory a backend it has no implementation for. Both DSAMLAConfig/NoPEDSAMLAConfig fields now default to tilelang directly and explicitly.
Bound the KPool indexer's logits tile with query chunking, mirroring the existing tilelang DSA selector's query-chunk bound.

Known gaps: tilelang sparse_mla_backend and deep_gemm_fp8 indexer_backend are not implemented for NoPE (explicit NotImplementedError, not silent fallback); KPool/NoPE-DSA sequence-parallel path is written per design but not yet GPU-tested at sp_size>1.

Test Plan

tests/model/test_glm53_dsa.py (12), test_glm53_nope_dsa_mla.py (2), test_flash_mla_cudnn_sparse_mla.py (3) all pass; GLM-5.2's existing tests/module/attention/test_dsa_mla.py (15) rerun clean against the shared tilelang.py/dsa_mla.py edits, confirming no regression. Numerical oracle: transformers 5.17.0's glm5_next model.

jayhenry · 2026-10-09T06:52:52Z

        self.up_proj = build_linear(self.hidden_size, self.intermediate_size, bias=mlp_bias, float8_cfg=float8_cfg)
        self.down_proj = build_linear(self.intermediate_size, self.hidden_size, bias=mlp_bias, float8_cfg=float8_cfg)
-        self.act_fn = get_act_fn(hidden_act)
+        self.act_fn = get_gated_act_fn(hidden_act, swiglu_limit)


去掉 get_gated_act_fn，使用 MoEActFnConfig ，是不是更好的方案？

jayhenry · 2026-10-09T08:03:56Z

        from .flash_mla import flash_mla_sparse_mla

        return flash_mla_sparse_mla
+    if backend == "flash_mla_cudnn":


cudnn_dsa 已经改为前向 flash mla，去掉这个分支？

jayhenry · 2026-10-09T15:39:21Z

+        starts,
+        ends,
+        select_k,
+        query_chunk_size=query_chunk_size,


增加 selector 参数，以便支持 tilelang_deepselect

jayhenry · 2026-10-09T15:41:11Z

+def get_kpool_topk_indices(backend: KPoolIndexerBackend) -> KPoolTopKIndicesProtocol:
+    if backend == "torch":
+        return torch_kpool_topk_indices
+    if backend == "tilelang":


支持 tilelang_deepselect

…iGLU Implements the F5 milestone from doc/xtuner_glm5p3flash_design.md: the NoPE (qk_rope_head_dim=0) DeepSeek Sparse Attention layers, their KPool indexer (pools of index_kpool consecutive tokens scored together instead of per-token top-k), a new flash_mla_cudnn SparseMLA backend (FlashMLA forward + existing cuDNN backward), and the clamped SwiGLU activation used by GLM-5.3-Flash's dense/shared/MoE MLPs. - xtuner/v1/ops/act_fn.py: native_clamped_swiglu + MoEActFnConfig support. - xtuner/v1/module/decoder_layer/{dense,moe}_decoder_layer.py: wire swiglu_limit through DenseMLP/MoEMLP. - xtuner/v1/ops/sparse_mla/kpool.py: pool layout, causal ranges, and top-k pool selection (torch reference + TileLang-backed production path, reusing the existing indexer kernel unmodified since relu's homogeneity makes head_dim^-0.5 movable from the relu argument into the per-head weight without changing the result). - xtuner/v1/ops/sparse_mla/flash_mla_cudnn.py: new SparseMLA backend. - xtuner/v1/ops/sparse_mla/tilelang.py: widen the hardcoded 576 head-dim check to a (head_dim, value_dim) whitelist. - xtuner/v1/model/moe/glm52/dsa_mla.py: reject flash_mla_cudnn as an indexer_backend (SparseMLABackend widened, but this backend has no indexer counterpart). - xtuner/v1/model/moe/glm53/nope_dsa_mla.py: NoPEDSAMLAConfig, KPoolIndexer, NoPEDSAMultiLatentAttention. - xtuner/v1/ops/sparse_mla/protocol.py: KPoolIndexerBackend (restricts GLM-5.3-Flash's KPool indexer_backend to "torch"/"tilelang" -- unlike GLM-5.2's 6-way DSAIndexerBackend, there's no cudnn_dsa/flash_mla/ deep_gemm_fp8/cute_dsl KPool kernel) and KPoolTopKIndicesProtocol. - xtuner/v1/ops/sparse_mla/__init__.py: get_kpool_topk_indices(backend), mirroring get_dsa_topk_indices's style -- explicitly raises for any backend other than "torch"/"tilelang" instead of silently falling through. KPoolIndexer resolves this once in __init__ into self._topk_indices_fn instead of re-branching on every forward call. - xtuner/v1/model/moe/glm53/nope_dsa_mla.py: NoPEDSAMultiLatentAttention. forward branches on freeze_dsa_indexer (torch.no_grad() only when frozen), mirroring GLM-5.2's per-token indexer, ahead of a future differentiable indexer output -- today's kpool_topk_indices/ torch_kpool_topk_indices still only ever return an int32 index tensor, so this doesn't yet change what's trainable (confirmed via a freeze_dsa_indexer=False smoke run: indexer params get requires_grad=True but no actual gradient). Calls self.indexer(...) directly instead of through reuse_during_recompute, which retained topk_ids' activations across the backward recompute pass to avoid recomputing them; dropped since nothing currently offloads or reuses that retained memory. Known gaps recorded in doc/progress.md: tilelang sparse_mla_backend and deep_gemm_fp8 indexer_backend are not implemented for NoPE (explicit NotImplementedError, not silent fallback); KPool/NoPE-DSA sequence- parallel path is written per design but not yet GPU-tested at sp_size>1; whether index topk ids can be offloaded is still an open question, not attempted here. Test plan: tests/model/test_glm53_dsa.py (12), test_glm53_nope_dsa_mla.py (2), test_flash_mla_cudnn_sparse_mla.py (3) all pass; GLM-5.2's existing tests/module/attention/test_dsa_mla.py (15) rerun clean against the shared tilelang.py/dsa_mla.py edits, confirming no regression. Numerical oracle: transformers 5.17.0's glm5_next model. Two near-tied top-k floating-point sensitivities were root-caused via seed sweeps and a targeted trace script (not worked around by loosening tolerances blindly) -- see doc/progress.md F5 section for the full analysis. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Also folds in what were originally two follow-up commits, now part of F5 itself rather than separate history: - Give the DSA backends explicit defaults instead of inheriting: both DSAMLAConfig/NoPEDSAMLAConfig fields now default to tilelang directly, with indexer_backend never falling back to sparse_mla_backend (the two name different, only partially overlapping backend vocabularies). - Bound the KPool indexer's logits tile with query chunking, mirroring the existing tilelang DSA selector's query-chunk bound.

* [Fix] Define and forward the TileLang DSA indexer selector 3de92e7 introduced the DeepSelect selector switch but referenced `selector` inside tilelang_indexer_topk_from_ranges without declaring it in the signature, so every call raised NameError. GLM-5.3's default KPool indexer backend (indexer_backend="tilelang") routes through this wrapper, so training could not start a single step on the f5 stack. The outer tilelang_dsa_topk_indices also accepted a selector but dropped it when calling the wrapper, so the tilelang_deepselect backend silently ran the torch.topk kernel instead of DeepSelect's radix-select kernel. - Declare selector in tilelang_indexer_topk_from_ranges, defaulting to "torch" to preserve the pre-DeepSelect behavior for existing callers (the KPool indexer). - Forward selector from tilelang_dsa_topk_indices. - Add CPU routing tests for the wrapper plus a CUDA-gated forwarding test that would have caught both regressions. * [Fix] Guard NoPE DSA config against unsupported settings at config time Two config-time guards for NoPEDSAMLAConfig, closing review findings #1 and #3 from the GLM-5.3-Flash stack review: - freeze_dsa_indexer=False now raises ValueError in build(), mirroring GLM-5.2's guard. The indexer only returns int32 indices and its gradients are cut by no_grad, so the flag silently left indexer params requires_grad=True -- paying optimizer state and activation memory for parameters that never train, with no error. - sparse_mla_backend='flash_mla_cudnn' (the production default) now requires num_attention_heads % 64 == 0 in the model validator, instead of failing at the first forward inside the FlashMLA kernel. The production config (64 heads) is unaffected; the torch reference backend keeps no alignment requirement. Also adds three tests covering both rejections and the torch-backend exemption. --------- Co-authored-by: liutong <liutong@pjlab.org.cn>

b03bf01/fe0037f0 switched _cudnn_dsa_sparse_mla_backward_op from the log2-space LSE contract to natural-log LSE, but this backend kept forwarding the log2 value, so every GLM-5.3 training silently fed the cuDNN backward a 1.4427x-inflated LSE. Forward is unaffected; backward gradients come out systematically shrunk (grad_norm ~20-30% low) and loss drifts (+0.043 over 300 steps vs the pre-change baseline). Save and pass the natural-log softmax_lse instead. Co-authored-by: liutong <liutong@pjlab.org.cn>

The tiny GLM-5.2 cases selected TileLang dimensions and sparse MLA kernels outside their supported shape set. Keep the actual indexer path while using a supported index dimension and the torch sparse MLA reference. The F5 routed activation check now uses the F5 public config rather than importing the F6 text model.

The colocate test inherited WORLD_SIZE from other cases and requested 32 workers on an 8-GPU node. Direct sampled KL is noisy for greedy rollouts, so retain its finite check and bound the stable K3 metric. Check temporary CUDA tensor lifetime directly instead of requiring allocator reserved bytes to decrease.

…2155) * [Fix] Compact KPool indices before cuDNN SparseMLA backward cuDNN consumes a valid prefix, while KPool indices can contain leading or interspersed -1 slots. Counting valid entries without moving them changes the key set used in backward and can produce exploding gradients. Stably compact valid indices before the cuDNN call, preserving key order and duplicates. Keep the FlashMLA forward and KPool selection unchanged. (cherry picked from commit e822121) * [Perf] Add cooperative BF16 KPool scoring and DeepSelect Add an independently selectable cooperative TileLang producer and DeepSelect adapter with reusable 1024-query FP32 scratch. Canonicalize pool ordering before expansion, preserve packed causal ranges/tails, and retain the original backend. Keep the original TileLang backend as the default; opt in with indexer_backend="tilelang_cooperative". Include focused regression tests. Previously measured in the tested source implementation, before this F5 port: H200, BF16 Q/K, 32 heads x 128, FP32 head weights, KPool size 4, top512 pools, query chunk 1024. Across all 131072 queries on one GPU, median CUDA-event time over five samples: - Scoring plus selection: original chunk1024 124.379 ms versus cooperative scoring plus DeepSelect 36.451 ms (3.41x faster). - Full indexer including projections: 174.623 ms versus 86.552 ms (2.02x faster). For the final 1024 causal queries at 128K, median CUDA-event time over 15 samples: scoring including original score cleanup 0.919424 ms versus 0.425504 ms (2.16x); top-k on identical logits 0.513184 ms versus 0.091872 ms (5.59x). Synthetic eager measurements include allocation/launch gaps and exclude compilation. Short 1K/4K full-indexer cases do not improve. These are prior source-workspace measurements, not a new GPU benchmark of this branch. Score comparison measured maximum absolute error 4.768e-7 and relative L2 error 6.561e-8. Selection may differ near the cutoff: 131065/131072 full-indexer rows were exact, with selected-ID overlap 99.9999895% at 128K. Preserve packed-document causal and tail semantics. * [Perf] Balance causal KPool scoring across SP ranks Redistribute scoring queries and FP32 head weights with neighbor P2P while preserving uniform attention ownership. Return compact pool IDs before expansion; cache aligned causal-work partitions and use collective fallback. Retain the cooperative scorer and DeepSelect backend, query chunk 1024, and the selected cuDNN backward tail fix. Keep SP balancing disabled by default; opt in with indexer_balance_sp=True. Previously measured in the tested source implementation, before this F5 port: eight H200 GPUs, SP8, batch 1, 131072 global tokens, BF16, KPool size 4, top512 pools, query chunk 1024, fixed-work ratio 0.5. Full cooperative indexer latency decreased from 14.992 ms unbalanced to 13.688 ms balanced: 8.70% less time, or 1.095x throughput. Two warmups and seven interleaved measurements per mode on identical features; take maximum-rank wall time per iteration, then mean. Redistribution and return communication are included; compilation/profiling are excluded. The complete four-layer block change was small relative to overlapping sample variation; no robust whole-block throughput improvement is claimed. Balancing is forward-only, so backward timing differences should not be attributed to it. The ratio-0 causal-only partition was slower (14.934 ms off versus 16.396 ms on); retain the measured ratio 0.5. These are prior source-workspace measurements, not a new GPU benchmark of this branch. Validate partition arithmetic, packed and poolless sequence indices, multi-hop transfers, and configuration wiring with CPU and four-rank Gloo regression tests. Prior SP8 validation preserved exact selected-token IDs and canonical ordering, including 128K. * [Perf] Compile KPool token-index expansion by default Compile expand_pools_and_tail by default with torch.compile. Both original TileLang and cooperative KPool selectors use the same integer expansion path; pooling and scoring remain as before. No new configuration option or kernel is added. At 128K tokens on one H200, pool size 4, top512 pools (2048 tokens), output alignment 512 and query chunk size 1024, compiling expansion alone reduces full-indexer CUDA-event time from 86.53 to 42.25 ms with cooperative scoring (2.05x, 51.2% reduction), and from 175.61 to 131.17 ms with the original scorer (1.34x). Measurements are medians of ten interleaved samples after two warmups per mode; cold compilation and autotuning are excluded. Complete indexer outputs are bitwise identical before and after the expansion-only change at 128K for both backends. Standalone expansion previously measured 46.62 to 2.43 ms (19.22x). Validation: 43 existing focused indexer, top-k, packed-layout, compile, sequence-parallel and cuDNN tests passed. SP8 selected-ID ordering, causal limits and tails are checked at 128K and on packed/poolless documents. Ruff, formatting and whitespace checks pass. --------- Co-authored-by: Knlife <2512955763@qq.com> Co-authored-by: shijiayang <sjyouuuuug@gmail.com>

jayhenry added this pull request to stack #2112 September 23, 2026 06:54

jayhenry force-pushed the feat/glm53flash-f5-nope-dsa branch from 8557cda to 9b26573 Compare October 8, 2026 10:12

This was referenced Oct 8, 2026

[Test] Check GLM-5.3 end-to-end logits alongside loss #2139

Closed

[Test] Run text and image cases through shared GLM parity checks #2140

Closed

jayhenry force-pushed the feat/glm53flash-f5-nope-dsa branch from 9b26573 to c2091a2 Compare October 9, 2026 06:32

jayhenry commented Oct 9, 2026

View reviewed changes

jayhenry force-pushed the feat/glm53flash-f5-nope-dsa branch from eb6ad8f to 5344d1d Compare October 10, 2026 19:01

jayhenry force-pushed the feat/glm53flash-f5-nope-dsa branch 2 times, most recently from 0cc9503 to 1199a3b Compare October 14, 2026 16:25

jayhenry and others added 7 commits October 14, 2026 17:14

[Chore] Make the TileLang scorer pass lint

3196895

jayhenry force-pushed the feat/glm53flash-f5-nope-dsa branch from 1199a3b to 3196895 Compare October 14, 2026 17:17

Navigation Menu

Repository navigation

[Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU - #2108

[Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU#2108
jayhenry wants to merge 7 commits into
feat/glm53flash-f4-mhcfrom
feat/glm53flash-f5-nope-dsa

jayhenry commented Sep 23, 2026 •

edited

Loading

Uh oh!

jayhenry Oct 9, 2026

Uh oh!

jayhenry Oct 9, 2026

Uh oh!

jayhenry Oct 9, 2026 •

edited

Loading

Uh oh!

jayhenry Oct 9, 2026

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Conversation

jayhenry commented Sep 23, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Test Plan

Uh oh!

jayhenry Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

jayhenry Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

jayhenry Oct 9, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

jayhenry Oct 9, 2026

Choose a reason for hiding this comment

Uh oh!

This branch has not been deployed

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

jayhenry commented Sep 23, 2026 •

edited

Loading

jayhenry Oct 9, 2026 •

edited

Loading