Repository navigation
Commit 8557cda
committed
[Feature] Add GLM-5.3-Flash F5: NoPE DSA + KPool indexer + clamped SwiGLU
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.1 parent 8ffc8e2 commit 8557cda
25 files changed
Lines changed: 2021 additions & 37 deletions
File tree
- doc
- examples/v1/config
- tests
- engine
- model
- module
- attention
- ops
- xtuner/v1
- data_proto
- model/moe
- glm52
- glm53
- module/decoder_layer
- ops
- sparse_mla
- utils
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
189 | 189 | | |
190 | 190 | | |
191 | 191 | | |
192 | | - | |
| 192 | + | |
| 193 | + | |
| 194 | + | |
| 195 | + | |
| 196 | + | |
| 197 | + | |
| 198 | + | |
| 199 | + | |
| 200 | + | |
| 201 | + | |
| 202 | + | |
| 203 | + | |
| 204 | + | |
| 205 | + | |
| 206 | + | |
| 207 | + | |
| 208 | + | |
| 209 | + | |
| 210 | + | |
| 211 | + | |
| 212 | + | |
| 213 | + | |
| 214 | + | |
| 215 | + | |
| 216 | + | |
| 217 | + | |
| 218 | + | |
| 219 | + | |
| 220 | + | |
| 221 | + | |
| 222 | + | |
| 223 | + | |
| 224 | + | |
| 225 | + | |
| 226 | + | |
| 227 | + | |
| 228 | + | |
| 229 | + | |
| 230 | + | |
| 231 | + | |
| 232 | + | |
| 233 | + | |
| 234 | + | |
| 235 | + | |
| 236 | + | |
| 237 | + | |
| 238 | + | |
| 239 | + | |
| 240 | + | |
| 241 | + | |
| 242 | + | |
| 243 | + | |
| 244 | + | |
| 245 | + | |
| 246 | + | |
| 247 | + | |
| 248 | + | |
| 249 | + | |
| 250 | + | |
| 251 | + | |
| 252 | + | |
| 253 | + | |
| 254 | + | |
| 255 | + | |
193 | 256 | | |
194 | 257 | | |
195 | 258 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
914 | 914 | | |
915 | 915 | | |
916 | 916 | | |
917 | | - | |
| 917 | + | |
918 | 918 | | |
919 | 919 | | |
920 | 920 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
59 | 59 | | |
60 | 60 | | |
61 | 61 | | |
62 | | - | |
63 | | - | |
64 | | - | |
65 | | - | |
| 62 | + | |
| 63 | + | |
| 64 | + | |
| 65 | + | |
66 | 66 | | |
67 | 67 | | |
68 | 68 | | |
| |||
| Original file line number | Diff line number | Diff line change | |
|---|---|---|---|
| |||
87 | 87 | | |
88 | 88 | | |
89 | 89 | | |
| 90 | + | |
90 | 91 | | |
91 | 92 | | |
92 | 93 | | |
| |||
0 commit comments