[megatron] GLM-5.3-Flash (glm5_next) support: KDA + NoPE-MLA/DSA + mHC hybrid MoE, 4-layer GPU CI entry

[erictang000]

What

Adds GLM-5.3-Flash (HF glm5_next, Glm5NextForConditionalGeneration) to the Megatron backend and a test_megatron_models H100 CI entry on the 4-layer slice eatang/GLM-5.3-Flash-4layer (mirror of CharyZeng/GLM-5.3-Flash-4layer; real truncated weights, so the logprob distribution stays peaked and the vLLM/Megatron comparison is meaningful even though the slice is not a coherent LM).

The architecture is unknown to both pinned upstreams: 45 layers of 3:1 KDA linear-attention / NoPE-MLA DeepSeek-sparse-attention (kpool indexer), 288-expert sigmoid MoE with one shared expert (dense MLP in the first layer), clamped SwiGLU, and Manifold-Constrained Hyper-Connections (mHC) on every block, shipped as a VL checkpoint.

Written from the HF modeling code against upstream Megatron-LM / Megatron-Bridge APIs, reusing what upstream already ships: megatron-core's DSAttention + AbsorbedMLASelfAttention (the NoPE case, qk_pos_emb_head_dim=0, works as is), TEGroupedMLP with activation_func_clamp_value for the clamped SwiGLU, and — since #2180 bumped megatron-core to b3393bbb — megatron-core's own HyperConnectionModule and TransformerBlock mHC stream handling.

Components

Model-specific code lives in two packages, split by where it should eventually land upstream.

workers/megatron/mcore_ext/ → Megatron-LM (megatron.core)

file	what	upstream status
hyper_connection.py (49)	RMSNormInputHyperConnectionModule: megatron-core's HyperConnectionModule with the input normalization overridden to GLM's standard RMSNorm, x * rsqrt(mean(x^2) + rms_norm_eps), instead of upstream's x / (rms(x) + 1e-6).	Delete once TransformerConfig carries the input-norm knobs (mhc_norm_eps / mhc_norm_eps_inside_sqrt).
mhc_transformer_layer.py (199)	HyperConnectionTransformerLayer with MoE MLP support — megatron-core's own mHC layer raises NotImplementedError for MoE sub-layers (mHC + MoE is only reachable there by wrapping MoE as a HybridStack layer), which GLM-5.3-Flash needs on all but the first layer. TransformerBlock owns the block-boundary stream expand/contract; this layer implements only the per-sub-layer residual update. Rejects 'mhc' in recompute_modules, whose managers it does not thread.	The MoE support is the piece to upstream.
kda.py (327)	KimiDeltaAttention: KDA linear attention (Kimi Linear / GLM-5.3-Flash) from TE column/row/duplicated linears, three depthwise causal convs, fp32 A_log/dt_bias, fla chunk_kda with the in-kernel safe gate (lower_bound * sigmoid(exp(A_log) * (f + dt_bias))) and fla FusedRMSNormGated output gate. TP shards heads (conv/A_log/dt_bias along heads; low-rank f_a/g_a duplicated; replicated o_norm grads summed across TP); packed thd via cu_seqlens; sharded_state_dict like GDN. Reuses the GDN config fields plus kda_gate_lower_bound. No CP, no inference cache.	New experimental_attention_variant="kda" candidate next to gdn/gdn2.

workers/megatron/glm5_next/ → Megatron-Bridge (models/glm/)

file	what
provider.py (55)	Glm5NextModelProvider(MLAModelProvider): mHC fields named after megatron-core's TransformerConfig, plus kda_gate_lower_bound, mhc_norm_eps* and dsa_indexer_kpool*. The attention pattern uses the generic linear_attention_freq list (1 = KDA, 0 = DSA).
layer_specs.py (113)	build_glm5_next_layer_spec: per-layer KDA-or-DSA × dense-or-MoE HyperConnectionTransformerLayer specs, built from the public get_dsa_module_spec_for_backend / get_moe_module_spec_for_backend / get_mlp_module_spec_for_backend helpers.
dsa.py (49)	Glm5NextDSAttention(DSAttention): the k-pool indexer selects every visible token whenever a sequence has at most index_topk tokens, so megatron-core's token-level indexer with dsa_indexer_topk=index_topk is exact in that regime and is reused; longer sequences raise instead of silently attending to a different subset.
bridge.py (318)	Glm5NextBridge registered for Glm5NextForConditionalGeneration → GPTModel, model_type="glm5_next": config mapping off text_config (skipping the head_dim=0 RoPE width) and 1:1 parameter mappings under model.language_model.* — KDA (separate q/k/v_conv1d, A_log, dt_bias), mHC (hc_*_fn/hc_*_base replicated, hc_*_scale [3] ↔ three alpha_* scalars via a small custom mapping, the shape Megatron-Bridge's DeepSeek-V4 bridge already uses). All 3050 language-model tensors map; no dropped conversion tasks. The vision tower is not bridged (language_model_only=True required).

SkyRL glue

• patches/megatron/patch_fa4_cute_import.py + workers/megatron/__init__.py: flash-attn 2.8.x's FA4 flash_attn.cute raises AttributeError against the cutlass DSL the pinned vLLM pulls in, which megatron-core's except (ImportError, PackageNotFoundError) probe does not catch, so import megatron.core aborts. The guard marks flash_attn.cute unavailable only when it is actually broken. Still required on torch 2.13. Delete once megatron-core's probe catches Exception or the pins agree.
• MegatronWorker.init_configs: megatron-core rejects mHC under recompute_granularity="full", which is what SkyRL's default gradient_checkpointing=True resolves to. Its suggested alternative (selective recompute with 'mhc' in recompute_modules) needs the mHC recompute managers the layer above does not thread, so this downgrades to selective recompute of the remaining modules with a log line.
• model_bridges.py: imports the bridge to register it.
• test_megatron_models.py: glm-5.3-flash-4layer_h100_tp2_ep4 (Megatron TP2 EP4 ETP1 → DP2, vLLM TP4 colocated, language_model_only, packed sequences, inference_only_init, max_num_seqs=512).
• test_glm5_next_modules.py (new, single GPU): RMSNormInputHyperConnectionModule vs HF Glm5NextTextHyperConnection at the model's activation scale, and KimiDeltaAttention on packed sequences vs per-sequence HF Glm5NextTextLinearAttention.
• Docs: supported_models.mdx, .claude/docs/backends/megatron.md.

Dependencies

• vLLM → per-commit wheel 0.28.1rc1.dev359+g98ed0856f ([Model] add GLM-5.3-Flash support vllm-project/vllm#53906, the GLM-5.3-Flash model). Built for CUDA 13 / torch 2.13, matching the torch pinned on main after [chore] Upgrade to torch 2.13 #2175, so its vendored vllm.third_party.deep_gemm extension loads and the DSA indexer (which hard-requires DeepGEMM) works with no separately installed deep_gemm. Verified on this branch: vllm.third_party.deep_gemm._C imports, vllm.utils.deep_gemm.has_deep_gemm() is True, no standalone deep_gemm present.
• flashinfer 0.6.16.post3 → 0.6.18 (the pinned vLLM requires exactly 0.6.18).
• No torch, TE, megatron-core or megatron-bridge changes — torch 2.13 comes from [chore] Upgrade to torch 2.13 #2175 and the megatron revs from [deps][megatron] bump megatron-bridge to 8e5f0e1 and megatron-core to b3393bbb #2180.

vLLM 0.28.1 fallout handled: vllm.entrypoints.openai.cli_args moved to vllm.entrypoints.launchers.cli_args (imported with a fallback); /inference/v1/generate, which SkyRL's generation client calls, is now gated behind VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1 (set before build_app); vLLM registers a native sharded_rdt weight-transfer engine, so the SkyRL shim no-ops and its unit test accepts either engine.

Validation

Megatron vs HF transformers 5.16.1 (sdpa, bf16) logits on the 4-layer slice, 82-token GSM8K prompt+answer, packed thd input (standalone torchrun parity script).

On this branch (torch 2.13, megatron-core b3393bbb), HF reference recomputed on the same toolchain:

config	KL(HF‖Meg)	argmax agree	mean |Δlogprob|
TP1	0.0082	91.5%	0.058

For scale, #2156 measured HF-bf16 vs HF-fp32 on this same slice at 0.0071 KL / 95.7% argmax, so this sits at the bf16 noise floor. The slice predicts near-randomly (HF logprob per target token ≈ −14.15), which makes argmax agreement the noisiest of the three statistics — most positions are near-ties that flip on tiny numeric differences while KL barely moves.

test_glm5_next_modules.py on the same tree: 2 passed — mHC exact to 4e-9 against the HF module, KDA within 0.5% of the reference mean on packed sequences.

test_logprobs_matching_roundtrip[glm-5.3-flash-4layer_h100_tp2_ep4] on Anyscale 4×H100:

vLLM logprobs     - mean: -4.675364, std: 0.380763
Megatron          - mean: -4.686937, std: 0.391731
logprob diff mean: 0.063767, std: 0.130800          (threshold 1e-1)
vLLM logprobs after sync - mean: -4.654904, std: 0.410720
vLLM logprob diff mean: 0.262453, std: 0.217895     (threshold 3e-1, two independent greedy generations)

and again in the H100 CI suite alongside the other model entries (logprob diff mean: 0.063314, vLLM logprob diff mean: 0.250699). Both of those runs predate the merge with main — they were on torch 2.11 with the megatron bump applied. The suites need re-running on this branch's tree; earlier TP2/EP2 parity (KL 0.0066, argmax 98.8%) is from that same pre-merge toolchain.

Known issue

test_logprobs_matching_roundtrip[qwen3.5-0.8b-dense_tp2] OOMs on 22 GiB L4 in the megatron_models suite: vLLM reserves 19.83 GiB for KV cache at the default gpu_memory_utilization=0.9, then flashinfer 0.6.18's top_k_top_p_sampling_from_logits needs a ~970 MiB vocab-wide softmax during sampler warm-up. Introduced by the flashinfer bump here — that entry has no gmu override and passes on main's 0.6.16.post3. Needs either a gmu override for small cards or VLLM_USE_FLASHINFER_SAMPLER=0; not reproducible on H100 (80 GiB hides it).

Limitations / follow-ups

• k-pool indexer: not implemented. Exact for sequences ≤ index_topk (2048) tokens; longer sequences raise (glm5_next/dsa.py). Needs a DSAttention hook letting an indexer own the top-k selection (pool scoring → expand to tokens → append the query's tail pool).
• Cross-layer DSA index sharing (indexer_types containing "shared"; the 4-layer slice is all "full") raises: megatron-core's dsa_indexer_topk_freq/skip_topk_offset arithmetic runs over all layers and would pick KDA layers as sources in this hybrid.
• KDA: no context parallelism, no inference cache. mHC: PP=1 only (megatron-core's own restriction), and no selective layernorm/mlp/mhc recompute.
• Training is wired (TE linears, TP grad-sync attributes, fp32 A_log/dt_bias/mHC params) but this PR exercises forward + weight sync only.

Upstreaming plan

Megatron-LM (megatron-core)

1. HyperConnectionModule: add mhc_norm_eps / mhc_norm_eps_inside_sqrt for a standard-RMSNorm input norm, with the fused-kernel path extended or guarded. Retires mcore_ext/hyper_connection.py.
2. HyperConnectionTransformerLayer: allow MoE MLP sub-layers (currently NotImplementedError). Retires mcore_ext/mhc_transformer_layer.py.
3. KimiDeltaAttention as experimental_attention_variant="kda" (module + spec next to gated_delta_net, kda_gate_lower_bound in TransformerConfig). Retires mcore_ext/kda.py.
4. DSA: k-pool indexer support and a per-indexer top-k selection hook; hybrid-aware index sharing.

Megatron-Bridge

1. models/glm/glm5_next_{provider,spec,bridge}.py = this PR's glm5_next/ package, once (1)–(3) exist upstream — the bridge only needs its layer/provider imports repointed.
2. Promote the DeepSeek-V4 _HCAlphaMapping pattern (here HyperConnectionScaleMapping) into param_mapping.py so both bridges share it.

Once megatron-core carries (1)–(3), mcore_ext/ is deleted and glm5_next/ imports from megatron.core; once Megatron-Bridge carries the bridge, glm5_next/ goes too and model_bridges.py loses one import.

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces support for the GLM-5.3-Flash (glm5_next) model in the Megatron backend, implementing custom Megatron-Core extensions such as Kimi Delta Attention (KDA), Manifold-Constrained Hyper-Connections (mHC), and a custom transformer layer that supports MoE. It also updates dependencies (including vLLM and FlashInfer) and includes compatibility fixes for vLLM >= 0.28.1. The review identified two critical runtime issues: an AttributeError in the DSA attention module due to an incorrect attribute name (self.index_topk instead of self.dsa_indexer_topk), and a TypeError in the mHC transformer layer caused by passing an unsupported padding_mask argument to the MLP forward call.

[gemini-code-assist]

The attribute self.index_topk does not exist on DSAttention or Glm5NextDSAttention. The correct attribute name in megatron-core is self.dsa_indexer_topk (or self.config.dsa_indexer_topk). Accessing self.index_topk will raise an AttributeError at runtime.

Suggested change

	if max_seqlen > self.index_topk:
	raise NotImplementedError(
	f"GLM-5.3-Flash sparse attention with sequences longer than dsa_indexer_topk="
	f"{self.index_topk} tokens (got {max_seqlen}) needs the k-pool indexer, which the "
	if max_seqlen > self.dsa_indexer_topk:
	raise NotImplementedError(
	f"GLM-5.3-Flash sparse attention with sequences longer than dsa_indexer_topk="
	f"{self.dsa_indexer_topk} tokens (got {max_seqlen}) needs the k-pool indexer, which the "

[gemini-code-assist]

Standard Megatron-Core MLP and MoE layers (such as GroupedMLP or TEGroupedMLP) do not accept padding_mask in their forward method. Passing padding_mask=moe_padding_mask will raise a TypeError at runtime. You should remove this keyword argument.

Suggested change

	mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output, padding_mask=moe_padding_mask)
	mlp_output_with_bias = self.mlp(pre_mlp_layernorm_output)

[greptile-apps]

Confidence Score: 4/5

The DeepGEMM platform coverage must be fixed or the unsupported environments explicitly rejected before merging because GLM-5.3-Flash cannot initialize there.

Findings

1.  DeepGEMM fallback cannot load ▶

Diagram

%%{init: {'theme': 'neutral'}}%%
flowchart LR
    HF[HF GLM-5.3-Flash checkpoint] --> Bridge[Glm5NextBridge]
    Bridge --> Provider[Glm5NextModelProvider]
    Provider --> Block[HyperConnection transformer block]
    Block --> KDA[KDA layers]
    Block --> DSA[NoPE-MLA / DSA layers]
    DSA --> DG[DeepGEMM indexer kernels]
    Block --> MoE[Dense / MoE MLP]
    Block --> MHC[mHC residual streams]
    Bridge --> Sync[Weight synchronization]
    Sync --> VLLM[vLLM inference engine]

Loading

[greptile-apps]

DeepGEMM fallback cannot load

When the Megatron extra is installed on Linux aarch64 or x86_64 Python 3.13, the environment marker excludes the Torch 2.11-compatible deep-gemm package, so vLLM falls back to its Torch 2.13-linked extension and GLM-5.3-Flash inference fails to initialize its required DSA indexer.