[megatron] GLM-5.3-Flash (glm5_next) support + 4-layer GPU CI entry

[erictang000]

What

Adds GLM-5.3-Flash (glm5_next) to the Megatron backend and a test_megatron_models CI entry for
the 4-layer slice CharyZeng/GLM-5.3-Flash-4layer. Neither the pinned megatron-core nor Megatron-Bridge know the
architecture (45 layers = 34 KDA linear-attention + 11 DSA sparse-MLA with a kpool-compressed indexer,
288-expert sigmoid MoE, mHC hyper-connections on every block, NoPE MLA, VL checkpoint), so everything
model-specific lives in a new package skyrl/backends/skyrl_train/workers/megatron/glm5_next/:

file	what
mhc.py	mHC hyper-connection site (fp32 mapping, Sinkhorn comb, stream collapse/expand) in the HF parameter layout
layer.py	Glm5NextTransformerLayer: megatron-core TransformerLayer with the n-stream mHC residual update; expands streams in layer 1 and mean-contracts in the last layer so GPTModel sees [s, b, C]
kda.py	TP-sharded KDA (fla chunk_kda / fused_kda_gate / ShortConvolution) built from TE linears
dsa.py + kpool_indexer.py	NoPE MLA with the kpool lightning indexer. Sequences ≤ index_topk (2048) are exact causal attention and run through TEDotProductAttention (packed thd, FA2/cuDNN); longer sequences use a torch indexer + masked SDPA (exact, O(L²), CI-oriented)
layer_specs.py / provider.py	per-layer KDA/DSA block spec on top of get_gpt_decoder_block_spec, Glm5NextModelProvider(MLAModelProvider)
bridge.py / hf_config.py	Glm5NextBridge (registered for Glm5NextForConditionalGeneration, model.language_model.* prefixes, 1:1 mappings — separate q/k/v_conv1d, hc_*_scale kept as [3]), plus an AutoConfig shim for transformers < 5.16

Done on top of upstream megatron-core: upstream's HyperConnectionTransformerLayer rejects MoE sub-layers
(mHC + MoE is only supported via HybridModel there), so the mHC layer is SkyRL-owned.

Other changes:

• workers/megatron/__init__.py + patches/megatron/patch_fa4_cute_import.py: flash-attn 2.8.x's FA4
  flash_attn.cute module is incompatible with the cutlass DSL that vLLM ≥ 0.28 pins
  (cutlass.cute.core.ThrMma missing → AttributeError, which megatron-core's except ImportError
  probe does not catch), so import megatron.bridge currently fails in the locked megatron env. The
  guard marks flash_attn.cute unavailable only when it is actually broken.
• test_megatron_models.py: glm-5.3-flash-4layer_h100_tp2_ep4 (TP2 EP4 ETP1 → DP2, vLLM TP4
  colocated, language_model_only, packed sequences, FA2, inference_only_init, max_num_seqs=512).
• CPU unit tests for the kpool indexer against a port of the HF Glm5NextTextIndexer.
• 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
  model). Main binaries are stable-libtorch-ABI so they load against our torch 2.11, as with 0.28.0.
• DeepGEMM: vLLM's GLM-5.3 DSA indexer hard-requires it, and vLLM's vendored deep_gemm._C is
  linked against torch 2.13 (undefined symbol: c10::ValueError…) under torch 2.11. Added deep-gemm
  built from deepseek-ai/DeepGEMM@8b1392b (vLLM's pinned commit) against torch 2.11 + CUDA 13.
  Hosted at https://github.com/erictang000/DeepGEMM/releases/tag/v2.6.1-8b1392b-cu13torch2.11 (same
  pattern as the fast-hadamard-transform wheels). Verified: vLLM logs DeepGEMM PDL enabled and the
  DSA indexer runs; without it has_deep_gemm() is False and the GLM-5.3 engine refuses to start.
• transformers 5.8.0 → 5.16.1 (native glm5_next config; also required by the vLLM commit,
  >=5.10.4). Same pins as [deps] bump transformers to 5.16.1 #2155, so this stacks cleanly on it.
• flashinfer 0.6.16.post3 → 0.6.18 (the vLLM commit pins it).
• No TE bump was needed: DSA is NoPE with q/k/v head dim 256, so FA2 covers it.

Validation

Megatron vs HF transformers (5.16.1, eager) logits on the 4-layer slice, 117-token GSM8K
prompt+response (scratchpad/parity.py, torchrun):

Megatron config	KL(HF‖Meg)	argmax agree	token-logprob |Δ|
TP1, vs HF bf16	0.0091	94.0%	0.067
TP2 EP2, vs HF bf16	0.0078	94.9%	0.064
PP2, vs HF bf16	0.0091	94.0%	0.067
TP1, vs HF fp32	0.0078	94.9%	0.043
HF bf16 vs HF fp32 (noise floor)	0.0071	95.7%	0.056
TP1, 2340 tokens (sparse kpool path), vs HF fp32	0.0090	91.3%	0.056
HF bf16 vs HF fp32, 2340 tokens	0.0091	91.1%	0.058

i.e. the Megatron model is at the bf16 noise floor, including the >2048-token sparse-indexer path
and under TP/EP/PP sharding. All 3050 language-model tensors of the checkpoint map (no dropped
conversion tasks).

SkyRL test_logprobs_matching_roundtrip[glm-5.3-flash-4layer_h100_tp2_ep4] on 4×B200 (same
mesh as the CI entry: Megatron TP2 EP4, vLLM TP4 colocated), with the new lock:

vLLM logprobs     - mean: -4.680098, std: 0.384223
Megatron          - mean: -4.695891, std: 0.392579
logprob diff mean: 0.059706, std: 0.126375          (threshold 1e-1)
vLLM logprobs after sync - mean: -4.652792, std: 0.416090
vLLM logprob diff mean: 0.255952, std: 0.219843     (threshold 3e-1, two independent greedy generations)
1 passed in 374s

Standalone vLLM (same wheel, torch 2.11 + DeepGEMM) vs HF-bf16 on the reference sequence: token-logprob
|Δ| mean 0.062 — i.e. vLLM and Megatron are each one bf16-noise-floor away from HF and ~0.06 from each other.

vLLM 0.28.0 → 0.28.1rc1.dev359 fallout fixed in this PR (all CPU suites green: -m "not vllm"
1568 passed, -m vllm 192 passed):

• vllm.entrypoints.openai.cli_args moved to vllm.entrypoints.launchers.cli_args (import with fallback).
• /inference/v1/generate (SkyRL's generation endpoint) is now gated behind
  VLLM_ENABLE_SCALE_OUT_ENDPOINTS=1; set in vllm_server_actor before build_app.
• vLLM now registers a native sharded_rdt weight-transfer engine; SkyRL's shim already no-ops in that
  case, the unit test was asserting the shim's class. Whether vLLM's native engine is wire-compatible with
  SkyRL's vendored trainer side was not checked here (the sharded_rdt backend is not used by this test).
• Known non-issue: vllm.entrypoints.openai.api_server is now a deprecated re-export shim (warning only).

Not run here: the other megatron_models / h100 entries under the new lock (transformers 5.16.1 +
vLLM dev wheel) — please let the label-gated GPU workflows run on this PR.

Notes / follow-ups

• The 4-layer slice is not a coherent LM (ppl ≈ 8e5 on real text for HF too — that is the truncation,
  not a bug), but because the weights are real the logprob distribution is peaked (std ≈ 4.6 vs ≈ 0.5
  for a random-init tiny model), which is what makes the vLLM/Megatron logprob comparison sensitive.
• vLLM's KDA triton kernels launch with num_seqs * 64 heads in CUDA grid dim y; max_num_seqs=1024
  overflows the 65535 limit → Triton Error [CUDA]: invalid argument during CUDA-graph capture
  (reproduced identically on torch 2.13 / triton 3.7.1, so it is not a torch-2.11 artifact). Worth an
  upstream vLLM issue; SkyRL caps max_num_seqs for this model.
• fused (cuDNN) attention was not exercised on this box (libcudart 12/13 clash in the local env);
  the CI entry uses FA2.
• Context parallelism is not supported for KDA / the kpool indexer.
• Training (backward) path is wired (TE linears, grad sync attrs for replicated mHC/norm params,
  frozen indexer) but only forward + weight-sync is exercised by this test.

🤖 Generated with Claude Code