Conversation
…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.
jayhenry
added this pull request to stack #2112
September 23, 2026 06:54
This was referenced Sep 23, 2026
This branch has not been deployed
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
1 participant
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.
Stack (bottom to top):
feat/glm53flash-materialize-full-f0→mainfeat/glm53flash-f3-kda→feat/glm53flash-materialize-full-f0feat/glm53flash-f4-mhc→feat/glm53flash-f3-kdafeat/glm53flash-f5-nope-dsa→feat/glm53flash-f4-mhc← you are herefeat/glm53flash-f1-vl-data→feat/glm53flash-f5-nope-dsafeat/glm53flash-f2-vision-tower→feat/glm53flash-f1-vl-datafeat/glm53flash-f6-text-moe→feat/glm53flash-f2-vision-towerSummary
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 ofindex_kpoolconsecutive tokens scored together instead of per-token top-k), a newflash_mla_cudnnSparseMLA 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+MoEActFnConfigsupport, wired throughDenseMLP/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 supportstorch/tilelang, unlike GLM-5.2's 6-wayDSAIndexerBackend).Also folds in what were originally two separate follow-up commits, now part of this PR:
indexer_backendused to default toNoneand fall back tosparse_mla_backend, but the two name different, only partially-overlapping backend vocabularies (flash_mla_cudnnis SparseMLA-only;deep_gemm_fp8/cute_dslare indexer-only), so the fallback could hand the indexer factory a backend it has no implementation for. BothDSAMLAConfig/NoPEDSAMLAConfigfields now default totilelangdirectly and explicitly.Known gaps recorded in
doc/progress.md: tilelangsparse_mla_backendanddeep_gemm_fp8indexer_backendare not implemented for NoPE (explicitNotImplementedError, not silent fallback); KPool/NoPE-DSA sequence-parallel path is written per design but not yet GPU-tested atsp_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 existingtests/module/attention/test_dsa_mla.py(15) rerun clean against the sharedtilelang.py/dsa_mla.pyedits, confirming no regression. Numerical oracle: transformers 5.17.0'sglm5_nextmodel.