DeepSpeed

18e005aa - Add deepspeed.models.glm_moe_dsa: TileLang sparse attention for GLM-5.2 (#8659)

Add deepspeed.models.glm_moe_dsa: TileLang sparse attention for GLM-5.2 (#8659) # Add `deepspeed.models.glm_moe_dsa`: TileLang sparse attention for training GLM-5.2 **Part of:** #8672 (Q4 2026 roadmap), "Enable training of large scale models": "Test and verify support of SOTA models" (PoC @pengdurice). Sibling: #8658 (MiniMax-M3). ## Summary * GLM-5.2 cannot train past 16,384 tokens per GPU with the stock transformers attention. Its indexer scores and its attention mask are both dense `seq x seq` tensors. * This PR adds TileLang kernels that attend only to the 2,048 keys the indexer selects for each query, with a backward pass, and a replacement attention module that uses them. One call, `deepspeed.models.glm_moe_dsa.replace_attention(model)`, swaps it in. * With it, the full 78-layer GLM-5.2 (743.4B parameters) trains at 65,536 tokens per GPU on 8 nodes of 8 H200. On an 8-layer cut of the model at 16,384 tokens, 10 training steps give the same loss as the stock path and run 2.7 times faster. * It adds a new folder, `deepspeed/models/<family>/`, for training code that belongs to one model family. ## Why a `deepspeed/models/` folder Some models need changes that only make sense for that model: a kernel for its attention, a replacement module, a fix for its layout. Other training stacks keep this code in one folder per model family, for example `torchtitan/models/<model>/`, Megatron-Bridge's `src/megatron/bridge/models/<family>/`, and THUDM/slime's `slime_plugins/models/glm5/`. DeepSpeed has no such place on the training side. The inference side has `deepspeed/inference/v2/model_implementations/<family>/`, and AutoEP keeps one preset file per model in `deepspeed/module_inject/auto_ep_presets/`. The convention this PR starts: * One package per family in `deepspeed/models/`, named after the model's folder in `transformers.models`. Here that is `glm_moe_dsa`. A reader goes from a Hugging Face model to its DeepSpeed code without a lookup table. * A family package holds whatever that family needs: kernels, modules, or plain PyTorch. Here it holds the TileLang kernels in `kernels/` and the replacement module in `attention.py`. * A family package exposes the function to call on a built model before `deepspeed.initialize`. Here that is `replace_attention`. * `import deepspeed` does not import `deepspeed.models`. A family's optional dependency (here TileLang) is needed only by users of that family. The MiniMax-M3 sparse-attention PR adds a second family, `deepspeed/models/minimax_m3_vl/`. Both PRs add the same `deepspeed/models/__init__.py` and the same docs index entry, byte for byte, so they merge in either order without conflicts. ## Problem GLM-5.2 (transformers `glm_moe_dsa`, DeepSeek Sparse Attention) cannot train at long sequence lengths on the stock transformers path, whatever DeepSpeed does about parameters and optimizer state: * `GlmMoeDsaIndexer.forward` scores every query against every key as a dense float32 `[batch, 32 index heads, S, S]` tensor before the top-k. That is 34 GB at S = 16,384 and 137 GB at S = 32,768, for each layer, before the relu copy and the head-weighted sum. * `GlmMoeDsaAttention.forward` then turns the top-k indices into a dense `[batch, 1, S, S]` additive mask. It calls the attention interface with per-head K and V expanded from the 512-wide latent (`expand_kv`). So the attention is dense too, quadratic in S, and cannot use the absorbed form. * Registering another attention function through `AttentionInterface` does not help. By the time the interface is called, the module has already expanded K and V and built the mask. THUDM/slime trains GLM-5 in Megatron with TileLang kernels that do two things: 1. compute the indexer logits in blocks of query rows and take a hard top-k; 2. run attention in the absorbed multi-head latent attention form over the selected latent keys only, with a fused backward. Depending on slime or Megatron for these kernels is not an option, so this PR adapts them into DeepSpeed. ## The change **`deepspeed/models/__init__.py`** (new): the package docstring that states the convention above. It imports nothing. **`deepspeed/models/glm_moe_dsa/kernels/`** (new): the kernels. They are adapted from the `deepseek_v32` examples in tile-ai/tilelang (MIT), by way of the GLM-5 plugin in THUDM/slime (Apache-2.0). The kernel bodies are the same. The Megatron, context-parallel and sequence-parallel code is removed, and the latent width is a parameter instead of a constant. Importing this package imports TileLang. | file | what it holds | |---|---| | `sparse_mla_fwd.py` | forward kernel: online softmax over the gathered top-k latent keys, `-1` indices masked; returns the output and the log-sum-exp | | `sparse_mla_bwd.py` | backward: preprocess (delta), the `dq` kernel and the atomically accumulated `dkv` kernel, postprocess | | `sparse_mla.py` | `sparse_mla(q, kv, indices, scale, d_v)`: the `torch.autograd.Function` around the two | | `indexer_fwd.py` | lightning indexer logits for a block of query rows, `[rows, keys]` float32, `-inf` outside each row's window | | `indexer.py` | `select_topk(...)`: runs the indexer in blocks of 8,192 query rows (2 GiB of logits per block at 65,536 keys), takes the top-k per block, pads rows with fewer keys with `-1`; `@torch.no_grad()` | Shapes are token-major: `q [tokens, heads, 512 + 64]`, `kv [tokens, 1, 512 + 64]`, and `indices [tokens, 1, topk]` int32 with `-1` for unused slots. `topk` must be a multiple of 64. **`deepspeed/models/glm_moe_dsa/attention.py`** (new): `DeepSpeedGlmMoeDsaAttention(source)` and `replace_attention(model, index_block_rows=8192, warm_up=True)`. * The replacement keeps the source module's parameters and attribute names (`q_a_proj`, `q_a_layernorm`, `q_b_proj`, `kv_a_proj_with_mqa`, `kv_a_layernorm`, `kv_b_proj`, `o_proj`, `indexer`). State-dict keys, parameter-freezing patterns and AutoTP name patterns do not change. Only `forward` is different. * Forward, step by step: 1. The no-rope part of the query is absorbed into the latent with this head's slice of `kv_b_proj.weight`: `q_nope . (W_kc kv) == (W_kc^T q_nope) . kv`. 2. Rotary embedding is applied interleaved, as in the model. 3. The indexer selects the top-k positions. The scores are the same as in transformers: `sum_h w_h relu(scale q_h . k)`. The scale is folded into the weights, which is exact because relu is positively homogeneous. 4. `sparse_mla` attends over the selected latent keys. 5. The other slice of `kv_b_proj.weight` maps the output back to value space before `o_proj`. Layers of type `"shared"` reuse `prev_topk_indices` from the previous full layer, as the stock module does. * The module reads `kv_b_proj.weight` as a tensor and never calls `kv_b_proj`, so ZeRO-3 gets no forward hook for that weight. The module registers the weight with `deepspeed.zero.register_external_parameter`. ZeRO-3 then gathers it for this module's forward and backward and releases it afterwards. * The indexer is not trained. A hard top-k has no gradient, and transformers does the same (its indexer forward is `@torch.no_grad()`). Under ZeRO-3, users freeze the indexer parameters, as they already must on the stock path. * Training only: * no key-value cache; * no padding mask: every sequence in the batch is one full causal sequence; * attention dropout must be 0. * Replace every DSA attention layer of a model, or none. The stock module encodes the earliest rows' missing keys through its causal mask, and this one through `-1`, so a shared layer must receive indices from a layer of its own kind. * Constructor checks: * `kv_lora_rank` and `qk_rope_head_dim` are powers of two; * `index_topk % 64 == 0`; * the head count is at most 64 or a multiple of 64; * the index head count divides 128. * `replace_attention` checks for TileLang first and raises `ImportError` with the install command. * Kernel warm-up. With `warm_up=True` (the default), when the backward kernel uses 128 threads (head blocks under 64 heads, for example with the heads split over tensor-parallel ranks), each attention module launches every kernel once, forward and backward, on a 128-token input with the model's head counts, head dims and top-k, at its first forward and on that forward's device. It runs without autograd, so activation checkpointing's recomputation saves the same tensors as the first forward. * Why: the 128-thread backward kernel spills registers, so its first launch makes the CUDA driver reserve local memory outside PyTorch's allocator, and the driver keeps it. Without the warm-up that launch happens in the first backward of training, when the allocator's cache can hold the rest of the device. On the full 78-layer GLM-5.2 at 65,536 tokens per GPU on 64 H200s it failed there with `CUDA_ERROR_OUT_OF_MEMORY` on 3 of the 64 GPUs. * The 256-thread kernel (64-head blocks: GLM-5.2's 64 heads on one GPU) does not spill, so it gets no warm-up. Measured on one H200, 64 heads, 32,768 tokens, memory used outside PyTorch's allocator around the first backward launch: | Backward kernel | First training-length launch, no warm-up | 128-token warm-up launch | Training-length launch after the warm-up | |---|---|---|---| | 128 threads | +0.469 GiB | +0.467 GiB | +0.000 GiB | | 256 threads (64-head blocks) | +0.002 GiB | not run | | * `kernels/sparse_mla_bwd.py` `backward_threads(heads_per_kv_group)` picks the thread count for both the kernel and the warm-up decision, so the two cannot disagree. * The forward and indexer kernels take the sequence length at run time, so the warm-up builds the kernels training uses. The backward kernel takes the batch size and both sequence lengths at compile time, so training's first backward still compiles its own backward kernel; what the warm-up launch provides is the local-memory reservation, which covers the training-length launch (table). **`docs/code-docs/source/models/glm_moe_dsa.rst`** (new) and an entry in **`docs/code-docs/source/index.rst`** that lists every page in `models/`: usage, requirements, limits, and the API reference. **`tests/unit/v1/models/test_glm_moe_dsa.py`** (new): see Tests. ## Changes after review Rebased on master (`5f266e26d`, 2026-10-15). All review points of the first round: | Review point | Change | |---|---| | Kernels adapted from TileLang / THUDM carry only DeepSpeed's header | The three kernel files carry the upstream copyright (`Copyright (c) Tile-AI.`) and `SPDX-License-Identifier: Apache-2.0 AND MIT`, and point to a new TileLang section in `THIRD_PARTY_NOTICES.md` (files, TileLang source paths and revisions, the slime revision they came through, the full MIT text). `scripts/check-license.py` accepts the expression. `setup.py` now lists `THIRD_PARTY_NOTICES.md` in `license_files`; before, the notice was in neither the sdist nor the wheel (checked by building both). | | `output_attentions=True` returns `attentions=()` | Raises `NotImplementedError`: the kernels never form attention weights. | | Padding and packed-sequence masks are ignored | A padding mask, a packed-sequence mask or position ids that restart raise `NotImplementedError`. No mask and eager attention's full causal mask are accepted. The check waits for the GPU, so it runs once per model forward, not once per layer. | | Widths 1 and 2 pass validation | `kv_lora_rank` and `qk_rope_head_dim` must be powers of two of at least 4 (the backward kernel splits the width into groups of 4). | | The warm-up runs on GPU 0 in every process | It runs at the module's first forward, on the input's device. | | The warm-up's backward kernel is compiled for 128 tokens | Correct: see the kernel warm-up bullet above. The warm-up now runs only for the 128-thread backward kernel, the one whose first launch reserves local memory. | Two additions from running GLM-5.2 at scale: * **256-thread backward kernel for 64-head blocks** (128 threads otherwise, as in tile-ai/tilelang#2592): 3.09x faster per kernel at GLM-5.2's shapes on an H200, 34-40% more tokens per second in training. * **Sequence split (context parallel).** `deepspeed.models.glm_moe_dsa.set_context_parallel_group(group)`: every rank of the group holds a contiguous slice of each sequence with global positions; the latent keys and the indexer keys are all-gathered over the group and their gradients reduce-scattered back; the indexer's causal bounds use global positions. GLM-5.2's first 38 layers at 65,536 tokens on 4 nodes x 8 H200, each sequence over 2 GPUs, ZeRO-3 + AutoEP: 96.4 s per step, 9.44% MFU, losses within 0.03 and the step-1 gradient norm within 0.02% of the same run on an earlier tree. ## Measured The at-scale runs below used this code before it moved into `deepspeed/models/`. The move changed only import paths and the class and function names. The kernel files are byte-identical. The unit tests were rerun on this branch (see Tests). ### Real weights, against Megatron-LM, 4 and 8 nodes x 8 H200 (2026-10-14 and 2026-10-15) These runs used this branch's files unchanged: the DeepSpeed tree they ran on has the same `deepspeed/models/glm_moe_dsa/` files. On one node (first 4 layers), master plus this PR gives the same loss at every step as that tree. Setup: GLM-5.2 real checkpoint, PG-19, 65,536 tokens per sequence, ZeRO-3 + AutoEP (expert parallel 8), each sequence over 2 GPUs (`set_context_parallel_group`), bf16 compute with fp32 master weights and Adam states, 80% (38 layers) or 95% (78 layers) of the optimizer on the CPU, full activation recompute. Megatron-LM: Megatron-FSDP, context parallel 4, expert parallel 8, the same share of the optimizer on the CPU. | Model | Nodes | Sequences per step | DeepSpeed, s/step | Megatron-LM, s/step | How measured | |---|---|---|---|---|---| | first 38 layers (348B parameters) | 4 x 8 H200 | 32 | 62.03 | 63.73 | one job on the same nodes, the frameworks alternating, mean of steps 4-23 of 2 runs each | | all 78 layers (743B parameters) | 8 x 8 H200 | 64 | 130.58 | 135.56 | separate jobs on 7 of the same 8 nodes, mean of steps 6-10 | Mean loss over the 23 steps of the same-job runs (first run of each): 9.6060 DeepSpeed, 9.6207 Megatron-LM. The DeepSpeed runs also used the ZeRO-3 pinned copy-back and fused CPU Adam changes (separate PRs). ### Full GLM-5.2, all 78 layers, 8 nodes x 8 H200, 65,536 tokens per GPU Setup: * Model: config-only build of GLM-5.2 with all 78 layers. 743.4B parameters, 724.8B of them in 75 MoE layers of 256 experts (top-8). 21 layers run the indexer, 57 reuse its selection. * Parallelism: AutoEP expert parallel 16 (4 expert groups of 2 nodes), TP 1, ZeRO-3 with `leaf_module` on the AutoEP layer. * Memory: optimizer state offloaded to host memory. Activation checkpointing with the checkpointed layer inputs in host memory (`checkpoint_in_cpu`). bf16 master weights and optimizer states. * Data: micro-batch 1, which is one 65,536-token PG-19 book per GPU. Tiled loss (16 shards). Indexer frozen. * Steps: 1 warm-up + 3 measured, learning rate 1e-6. * Environment: `NCCL_RUNTIME_CONNECT=0` and `NCCL_NVLS_ENABLE=0`. With NCCL 2.27's default lazy connection, the first scalar all-reduce of the optimizer's gradient norm needed new NCCL buffers after the backward had filled the device, and it failed. This is a launch setting, not part of this PR. * Versions: torch 2.8.0+cu126, transformers 5.15.1, tilelang 0.1.14, NCCL 2.27.3. | measured | value | |---|---| | step time, mean of 3 measured steps | 252.56 s (251.3 / 253.7 / 252.7) | | throughput | 16,607 tokens/s over 64 GPUs (4,194,304 tokens per step) | | MFU / HFU, H200 dense bf16 peak 989 TFLOPS | **7.28 %** / 10.98 % (72.0 TFLOPS per GPU of model FLOPs) | | peak allocated, highest GPU | 116.54 GiB of 139.80 | | peak reserved, highest GPU | 132.34 GiB | | allocated between steps (weights, gradients, buffers) | 33.15 GiB | | loss / gradient norm (random initialisation) | 13.19 / 4.62 to 4.97 | Same setup with `OMP_NUM_THREADS=12` for the CPU Adam optimizer, and DeepSpeed's wall-clock breakdown timers on. The launcher sets 1 thread per rank by default, which leaves the offloaded optimizer on one CPU thread. This is a launch setting, not part of this PR. | measured | value | |---|---| | step time | 220.5 s: forward 39.2 s, backward 165.9 s, optimizer step 14.8 s | | throughput | 19,020 tokens/s over 64 GPUs | | MFU | 8.34 % | How MFU is counted, per token: | part | GFLOP per token | |---|---| | 6 x the 40.10B active matmul parameters (attention projections, dense MLPs, router, shared expert, 8 routed experts, output layer) | 240.6 | | sparse attention over the selected keys (2,016 on average at this length), model head dims, forward + backward | 30.9 | | the frozen indexer, forward only | 6.1 | | **total** | **277.6** | HFU adds the forward recompute of activation checkpointing and the absorbed attention form that the kernel executes: 418.7 GFLOP per token. The stock transformers path cannot run this configuration: its indexer alone allocates a 137 GB score tensor per layer at 32,768 tokens. ### GLM-5.2 cut to 8 layers, 2 nodes x 8 H200, ZeRO-3 + AutoEP Setup: * Model: config-only build of GLM-5.2 cut to 8 of 78 layers. 5 MoE layers with 256 experts, top-8. The indexer schedule keeps both modes: layers 0, 1, 2 and 6 run the indexer; layers 3, 4, 5 and 7 reuse the selection. 52.5B parameters, of which 48.3B are expert weights. * Parallelism and memory: expert parallel 16, TP 1, ZeRO-3 with `leaf_module` on the AutoEP layer, bf16 master weights and optimizer states, Hugging Face gradient checkpointing. * Data: micro-batch 1 per GPU. Tiled loss (16 shards) for the 154,880-token vocabulary. Real text: PG-19 books re-tokenized with the GLM-5.2 tokenizer, one book per row. Indexer frozen. * Steps: 10 (3 warm-up + 7 measured). * Versions: torch 2.8.0+cu126, transformers 5.15.1, tilelang 0.1.14. * Throughput is tokens per second over all 16 GPUs. | attention | seq | s / step | tokens / s (16 GPUs) | step peak allocated | peak reserved | |---|---|---|---|---|---| | DSA kernels | 16,384 | 4.65 | 56,429 | 54.6 GiB | 67.0 GiB | | DSA kernels | 32,768 | 9.26 | 56,621 | 73.6 GiB | 91.5 GiB | | DSA kernels | 65,536 | 18.46 | 56,798 | 112.0 GiB | 133.4 GiB | | stock SDPA (memory-efficient kernel, dense top-k mask), second run | 16,384 | 12.70 | 20,645 | 71.6 GiB | 93.7 GiB | | DSA kernels, second run | 16,384 | 4.73 | 55,477 | 55.8 GiB | 68.2 GiB | * At 16K the kernels are 2.7 times faster than the stock path and reserve 25 GiB less. * Tokens per second stay flat from 16K to 65K. The selection is fixed at 2,048 keys per query, so attention cost grows linearly with the token count, like the rest of the layer. * No stock number exists past 16K on a 141 GB card. At 32K the stock indexer needs 137 GB for the scores of one layer alone. * The first step at 16K took 148 s, which includes the TileLang compile of both kernels for the run's shapes. The 32K and 65K runs started at 59 s and 70 s: ZeRO-3 first-step setup only, with the compile cached on the node. ### Stock transformers attention against the DSA module, same data, seq 16,384 The same 8-layer model, same seed, same PG-19 rows, learning rate 1e-5. The printed loss is the running mean of the bf16 per-step losses since step 1. | step | stock: running mean loss | DSA: running mean loss | stock: grad norm | DSA: grad norm | |---|---|---|---|---| | 1 | 13.250000 | 13.250000 | 8.2648 | 8.2634 | | 2 | 13.093750 | 13.093750 | 7.8671 | 7.8675 | | 3 | 12.916667 | 12.916667 | 8.5182 | 8.5209 | | 4 | 12.796875 | 12.796875 | 8.7343 | 8.7375 | | 5 | 12.675000 | 12.675000 | 8.3782 | 8.3754 | | 6 | 12.531250 | 12.531250 | 8.3926 | 8.3821 | | 7 | 12.491071 | 12.491071 | 9.5111 | 9.5051 | | 8 | 12.335938 | 12.335938 | 7.7833 | 7.7907 | | 9 | 12.319444 | 12.319444 | 6.6761 | 6.6875 | | 10 | 12.206250 | 12.206250 | 8.8822 | 8.8952 | * Every running-mean loss agrees to the printed precision, so every per-step bf16 loss is identical. * The gradient norms agree within 0.16 %. * Other learning rates: at 1e-6, ten steps do not move the loss. At 1e-4 this randomly initialised model diverges after one step, which amplifies bf16 noise and top-k ties. Even there, the per-step losses were identical for the first three steps. ## Tests `tests/unit/v1/models/test_glm_moe_dsa.py` (new). It is skipped without TileLang or a CUDA device. | test | what it checks, and the wrong behaviour it would catch | |---|---| | `test_sparse_mla_matches_reference` (8 heads, 64 heads) | the output and both input gradients of `sparse_mla`, against the sparse softmax written out in float32, over a random causal selection with `-1` padding. 8 heads exercises the kernel's padding of heads to 16 | | `test_indexer_logits_and_selection_match_reference` (8 and 32 index heads) | the logits against the float32 formula, with `-inf` exactly on the future keys. `select_topk` returns the reference top-k on every row where the k-th and (k+1)-th scores are apart, the right number of valid slots on the early rows, and never a future key | | `test_sparse_mla_at_glm52_sizes` | the run's shapes: 64 heads, top-k 2048 (32 index blocks of 64 per row), 16,384 tokens. Output, `dq` and `dkv` against the float32 reference on 256 random query rows. A row's output and `dq` depend on that row alone; `dkv` is checked by sending a gradient into those rows only | | `test_indexer_selection_at_glm52_sizes` | 32 index heads, 16,384 tokens (two 8,192-row blocks in `select_topk`), top-k 2048: exactly k distinct causal keys per row, each scoring at least the k-th best float32 score minus a rounding allowance, so no better key was left out | | `test_replacement_matches_transformers_module` | on a small `GlmMoeDsaConfig`, the replacement's selection overlaps the stock indexer's on more than 90 % of slots (the two differ at ties). On a shared layer given the same indices, the output, the input gradient and the gradient of every parameter (`q_a_proj`, `q_a_layernorm`, `q_b_proj`, `kv_a_proj_with_mqa`, `kv_a_layernorm`, `kv_b_proj`, `o_proj`) match the stock module on every row with a complete selection | | `test_replace_walks_the_model_and_keeps_parameter_names` | `replace_attention` swaps every layer of `GlmMoeDsaForCausalLM`, and the parameter names do not change | | `TestDsaAttentionZeroStage3` (2 ranks, ZeRO-3, `torch_adam`) | the loss and the gradient of `kv_b_proj.weight`, the weight read directly, match a single-process run of the same replaced model. Without the external-parameter registration the weight is a 0-element placeholder in the forward | | `test_unsupported_inputs_raise` | `output_attentions=True`, a padding mask and position ids that restart all raise `NotImplementedError` | | `test_kernel_widths_below_four_are_rejected` | `kv_lora_rank` or `qk_rope_head_dim` below 4 raises `ValueError` | | `test_first_forward_under_activation_checkpointing` | the warm-up at the first forward leaves nothing saved for backward, so non-reentrant activation checkpointing's recomputation does not fail | | `test_warm_up_runs_on_the_input_device` (2 GPUs) | with the current device still GPU 0 at `replace_attention`, the warm-up allocates nothing on GPU 0 when the model runs on GPU 1 | | `TestDsaAttentionContextParallel` (2 ranks) | each sequence split over 2 ranks gives the logits, the summed loss and every parameter gradient of the whole sequence on one rank (`index_topk` = sequence length, so the selection has no near-ties) | Results on this branch (master `5f266e26d` + this PR, 2026-10-15): **14 passed** in 167 s on one node of 8 H200 (torch 2.8.0+cu126, transformers 5.15.1, tilelang 0.1.14). Tests of the attention alone build the model with dense MLPs (`_tiny_config(dense_mlps=True)`): the MoE backward can hang in torch 2.8's `_grouped_mm` when an expert receives no tokens. `pre-commit run --files` passes on every changed file. ## Duplicate check Open and recent pull requests searched on 2026-10-15 for "glm", "sparse attention", "DSA", "TileLang", "indexer": none adds GLM-5.2 or DeepSeek-sparse-attention training kernels. #8658 is the sibling for MiniMax-M3 (another family in the same `deepspeed/models/` folder; both add the same `deepspeed/models/__init__.py` and docs entry, byte for byte). ## Compatibility | Cell | Validated | How | |---|---|---| | ZeRO-3 + AutoEP, bf16, CPU optimizer offload, activation checkpointing | yes | full 78-layer GLM-5.2, 8 nodes x 8 H200, 65,536 tokens per GPU; 38-layer cut on 4 nodes | | Sequence split (context parallel) over 2 GPUs | yes | `TestDsaAttentionContextParallel`; 38 and 78 layers at 65,536 tokens | | ZeRO-3, 2 ranks, plain DeepSpeed (no AutoEP) | yes | `TestDsaAttentionZeroStage3` | | Heads split over tensor-parallel ranks (128-thread backward) | unit tests only | warm-up tests; no at-scale run | | ZeRO-1 / ZeRO-2, fp16, pipeline parallel | no | | | Accelerators other than CUDA | no | TileLang kernels; `replace_attention` raises `ImportError` without TileLang | Later measurement (not part of this PR's diff): the full 78-layer model at 65,536 tokens per GPU on 8 nodes x 8 H200 runs at 130.6 s per step (MFU 14.1 %) with this PR's kernels together with other DeepSpeed changes (prefetch settings, a pinned copy-back of offloaded master weights, a fused CPU Adam step, 5 % of the optimizer on the GPU). The 252.56 s above is the 128-thread backward kernel with bf16 optimizer states. ## AI assistance This change, its tests and this description were prepared with AI assistance (Claude). Signed-off-by: pengdurice <pengduhit@gmail.com>