[megatron] GLM-5.3-Flash (glm5_next) support: KDA + k-pool DSA + mHC MoE, full fine-tuning and LoRA sync to vLLM (merge_lora=false)

.claude/docs/backends/megatron.md

@@ -8,6 +8,13 @@ SkyRL uses Megatron-Bridge for HF-to-Megatron model conversion. Installed from g
 - `MegatronConfig` in `skyrl/train/config.py`
 - `MegatronWorker` in `skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py`.
 - Custom bridges in `skyrl/backends/skyrl_train/workers/megatron/model_bridges.py` (e.g., `GLM47FlashBridge`).
+- GLM-5.3-Flash (`glm5_next`): `skyrl/backends/skyrl_train/workers/megatron/glm5_next/` (provider, block
+  spec, bridge; Megatron-Bridge-shaped) on top of `skyrl/backends/skyrl_train/workers/megatron/mcore_ext/`
+  (Megatron-Core-shaped: `hyper_connection.py` backport of Megatron-LM main's mHC module,
+  `mhc_transformer_layer.py` mHC layer with MoE support, `kda.py` KDA linear attention). The DSA layers reuse
+  megatron-core's `DSAttention` (NoPE MLA, token-level indexer) and are exact for sequences up to `index_topk`
+  (2048) tokens; the k-pool indexer for longer sequences is not implemented and raises. Requires
+  `language_model_only=True` (VL checkpoint) and packed sequences (`remove_microbatch_padding=True`).
 
 ## Parallelism Strategies
 

docs/content/docs/getting-started/supported_models.mdx

@@ -18,6 +18,7 @@ The following models are fully supported in SkyRL's Megatron backend via Megatro
 - Qwen-3.5 (Dense: [0.8B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_qwen3.5.sh)/2B/4B/9B/27B, MoE: [35B-A3B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_qwen3.5_35b_a3b.sh)/122B-A10B/397B-A17B)
 - Nemotron-3 (Dense: [Nano-4B-BF16](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/nemotron_3/run_nemotron_3_nano_4b_gsm8k.sh), MoE: Nano-30B-A3B-BF16)
 - [GLM-4.7-Flash](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_grpo_glm4_7_30b.sh)/GLM-4.7
+- GLM-5.3-Flash (language model; KDA + DSA hybrid MoE with mHC residuals, sequences up to `index_topk` tokens)
 - Qwen3 (Dense: [0.6B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron.sh)/1.7B/4B/8B/32B, MoE: [30B-A3B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_dapo_qwen3_30b_a3b.sh)/[235B-A22B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_qwen3-235b-a22b.sh))
 - [Moonlight-16B-A3B](https://github.com/NovaSky-AI/SkyRL/blob/main/examples/train/megatron/run_megatron_moonlight.sh)
 

pyproject.toml

@@ -115,8 +115,9 @@ skyrl-train = [
 
 fsdp = [
     "skyrl[skyrl-train]",
-    "vllm==0.28.0; sys_platform == 'linux'",
+    "vllm; sys_platform == 'linux'",
     "vllm-router; sys_platform == 'linux'",
+    "deep-gemm; sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'",
     # The `nixl` shim provides that namespace and dispatches on `torch.version.cuda`,
     # so with a cu13 torch it loads `nixl_cu13`. `nixl-cu13` ships the `nixl_cu13`
     # module, but vLLM imports `nixl._api`. Its metadata hard-depends on `nixl-cu12`
@@ -127,9 +128,9 @@ fsdp = [
     "causal-conv1d==1.6.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'",
     "flash-attn==2.8.3+cu.13.0.torch.2.11; sys_platform == 'linux'",
     "torch==2.11.0; sys_platform == 'linux'",
-    "flashinfer-python==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-jit-cache==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-cubin==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-python==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-jit-cache==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-cubin==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "torchvision; sys_platform == 'linux'",
 ]
 
@@ -151,23 +152,24 @@ megatron = [
     "flash-linear-attention; sys_platform == 'linux'",
     "causal-conv1d==1.6.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'",
     "mamba-ssm==2.3.2.post1+cu.13.0.torch.2.11; sys_platform == 'linux'",
-    "vllm==0.28.0; sys_platform == 'linux'",
+    "vllm; sys_platform == 'linux'",
     "vllm-router; sys_platform == 'linux'",
+    "deep-gemm; sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'",
     # The `nixl` shim provides that namespace and dispatches on `torch.version.cuda`,
     # so with a cu13 torch it loads `nixl_cu13`. `nixl-cu13` ships the `nixl_cu13`
     # module, but vLLM imports `nixl._api`. Its metadata hard-depends on `nixl-cu12`
     # too; that variant is overridden out below (it would drag the CUDA-12 runtime
     # into the image for a module the shim would never load).
     "nixl; sys_platform == 'linux'",
     "torch==2.11.0; sys_platform == 'linux'",
-    "flashinfer-python==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-python==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "torchvision; sys_platform == 'linux'",
     # megatron-bridge requires Python 3.12+; pin megatron-core to the same
     # constraint so both packages are consistently available (or absent).
     "megatron-bridge; sys_platform == 'linux' and python_version >= '3.12'",
     "megatron-core; sys_platform == 'linux' and python_version >= '3.12'",
-    "flashinfer-jit-cache==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-cubin==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-jit-cache==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-cubin==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     "nvidia-modelopt; sys_platform == 'linux'",
     "fast-hadamard-transform; sys_platform == 'linux'",
 ]
@@ -258,9 +260,9 @@ override-dependencies = [
     # so it doesn't pull a second, unused CUDA runtime in alongside the cu13 one.
     "nixl-cu12; sys_platform == 'never'",
     # Megatron-Bridge pins flashinfer-python==0.6.8.post1, which conflicts with
-    # our pin, so all three flashinfer packages are held here. vLLM 0.28.0 needs
-    # >=0.6.16: it imports `flashinfer.autotuner.set_autotune_process_group`,
-    # which 0.6.14 does not have, and the engine fails to start without it.
+    # our pin, so all three flashinfer packages are held here. The pinned vLLM
+    # (0.28.1rc1.dev359, see `vllm` in [tool.uv.sources]) requires exactly 0.6.18; vLLM
+    # 0.28.0 already needed >=0.6.16 (`flashinfer.autotuner.set_autotune_process_group`).
     # flashinfer hard-errors when an installed cubin's version differs from its
     # own, so the three must move together.
     # They must also be >=0.6.13: older flashinfer rejects the `layout_code` vLLM's
@@ -270,9 +272,9 @@ override-dependencies = [
     # flashinfer-cubin is not published to PyPI past 0.6.13, so it (like
     # flashinfer-jit-cache) comes from flashinfer's own index -- see the
     # `flashinfer` index and [tool.uv.sources] below.
-    "flashinfer-python==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-jit-cache==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
-    "flashinfer-cubin==0.6.16.post3; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-python==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-jit-cache==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
+    "flashinfer-cubin==0.6.18; sys_platform == 'linux' and platform_machine == 'x86_64'",
     # mamba-ssm 2.3.2.post1 pins tilelang==0.1.8 and apache-tvm-ffi<=0.1.9, while
     # vLLM 0.28.0 pins tilelang==0.1.12 and apache-tvm-ffi==0.1.11. Both overrides
     # hold these at the versions the lock already resolved to before mamba-ssm moved
@@ -292,9 +294,10 @@ override-dependencies = [
     # `_moe_C_stable_libtorch.abi3.so`), so its main binaries load against 2.11 and
     # we hold torch there rather than churn every CUDA extension pin. The one
     # exception is `vllm/third_party/deep_gemm/_C`, a version-specific cpython-312
-    # build linked against torch 2.13's c10: it fails to import and vLLM falls back,
-    # so the DeepGEMM fused-MoE / sparse-attn-indexer paths are unavailable. See
-    # `.claude/docs/weight_sync.md`. Revisit once those wheels publish 2.13 builds.
+    # build linked against torch 2.13's c10: it fails to import, so a `deep-gemm`
+    # wheel built against torch 2.11 is installed alongside (see [tool.uv.sources]);
+    # vLLM prefers it over the vendored copy. See `.claude/docs/weight_sync.md`.
+    # Revisit once those wheels publish 2.13 builds.
     "torch==2.11.0; sys_platform == 'linux'",
     # torchvision is NOT stable-ABI (unlike vLLM), and each release hard-pins one
     # torch: 0.28.0 wants torch 2.13.0, 0.26.0 wants 2.11.0. vLLM 0.28.0 pulls
@@ -399,6 +402,25 @@ megatron-bridge = {git = "https://github.com/NVIDIA-NeMo/Megatron-Bridge", rev =
 # `git ls-tree <bridge-rev> 3rdparty/Megatron-LM`. This one is on the Megatron-LM
 # main branch (https://github.com/NVIDIA/Megatron-LM/tree/main), dated 8/12/26.
 megatron-core = {git = "https://github.com/NVIDIA/Megatron-LM", rev = "14346b65a2d0790e451919858f7771078105c5f0", marker = "sys_platform == 'linux'"}
+# vLLM at 98ed0856f31fa3aaf5e27464e2b4ef5a8ee6b2f5 (vllm-project/vllm#53906, "[Model] add GLM-5.3-Flash
+# support", 2026-09-03), i.e. 0.28.1rc1.dev359 -- the first vLLM with the GLM-5.3-Flash (glm5_next: KDA +
+# DSA-kpool + mHC) model. Per-commit wheels are published by vLLM's CI at wheels.vllm.ai; this one is
+# built for CUDA 13 / torch 2.13 but, like 0.28.0, its main binaries are stable-libtorch-ABI so they load
+# against the torch 2.11 we hold (see the torch override above). Move back to a plain PyPI pin once a
+# release containing #53906 ships.
+vllm = [
+    { url = "https://wheels.vllm.ai/98ed0856f31fa3aaf5e27464e2b4ef5a8ee6b2f5/vllm-0.28.1rc1.dev359%2Bg98ed0856f-cp38-abi3-manylinux_2_28_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64'" },
+    { url = "https://wheels.vllm.ai/98ed0856f31fa3aaf5e27464e2b4ef5a8ee6b2f5/vllm-0.28.1rc1.dev359%2Bg98ed0856f-cp38-abi3-manylinux_2_28_aarch64.whl", marker = "sys_platform == 'linux' and platform_machine == 'aarch64'" },
+]
+# DeepGEMM at deepseek-ai/DeepGEMM@8b1392b (the commit the pinned vLLM vendors), built against torch 2.11 +
+# CUDA 13 (nvcc 13.0.88, cp312, x86_64, glibc >= 2.35) and hosted like the fast-hadamard-transform wheels
+# below. vLLM ships its own copy under `vllm.third_party.deep_gemm`, but that `_C` is a version-specific
+# cpython extension linked against torch 2.13 and fails to import under the torch 2.11 we hold; vLLM
+# prefers an installed `deep_gemm` package when one exists. Needed because vLLM's GLM-5.3-Flash DSA
+# indexer hard-requires DeepGEMM (DeepSeek-V3.2-style sparse-attention models too). Kernels are JIT
+# compiled at runtime, so the host still needs nvcc/nvrtc. Rebuild if the torch pin moves; drop once
+# vLLM's vendored extension loads against our torch.
+deep-gemm = { url = "https://github.com/erictang000/DeepGEMM/releases/download/v2.6.1-8b1392b-cu13torch2.11/deep_gemm-2.6.1%2B8b1392b-cp312-cp312-linux_x86_64.whl", marker = "sys_platform == 'linux' and platform_machine == 'x86_64' and python_full_version >= '3.12' and python_full_version < '3.13'" }
 # NOTE (sumanthrh): This custom wheel of vllm-router includes a fix for /chat/completions endpoint: https://github.com/vllm-project/router/pull/162  on top of 0.1.14
 # TODO (sumanthrh): Remove this after 0.1.15 vllm-router release
 vllm-router = [

skyrl/backends/skyrl_train/inference_servers/render_server.py

@@ -80,7 +80,11 @@ def _run_render_server(model_path: str, port: int, log_file: Optional[str]) -> N
         build_and_serve_renderer,
         setup_server,
     )
-    from vllm.entrypoints.openai.cli_args import make_arg_parser
+
+    try:  # vLLM >= 0.28.1 moved the CLI args to entrypoints.launchers
+        from vllm.entrypoints.launchers.cli_args import make_arg_parser
+    except ImportError:
+        from vllm.entrypoints.openai.cli_args import make_arg_parser
     from vllm.utils.argparse_utils import FlexibleArgumentParser
 
     async def _serve() -> None:

skyrl/backends/skyrl_train/inference_servers/utils.py

@@ -152,7 +152,11 @@ def build_vllm_cli_args(cfg: SkyRLTrainConfig) -> Namespace:
     """Build CLI args for vLLM server from config."""
     from vllm import AsyncEngineArgs
     from vllm.config import WeightTransferConfig
-    from vllm.entrypoints.openai.cli_args import FrontendArgs
+
+    try:  # vLLM >= 0.28.1 (entrypoints.launchers); the openai.cli_args path is the pre-0.28.1 location
+        from vllm.entrypoints.launchers.cli_args import FrontendArgs
+    except ImportError:
+        from vllm.entrypoints.openai.cli_args import FrontendArgs
     from vllm.platforms import current_platform
     from vllm.utils.argparse_utils import FlexibleArgumentParser
 

skyrl/backends/skyrl_train/inference_servers/vllm_server_actor.py

@@ -592,6 +592,9 @@ async def _build_and_serve_vllm_server(
     # One uvicorn per port (no api_server_count fan-out), matching vLLM's own
     # single-server path, so SO_REUSEPORT stays off.
     sock = create_server_socket(sock_addr, reuse_port=False)
+    # vLLM >= 0.28.1 gates the token-in/token-out ``/inference/v1/generate`` route (which SkyRL's
+    # generation client calls) behind this env var, read when the app is built.
+    os.environ.setdefault("VLLM_ENABLE_SCALE_OUT_ENDPOINTS", "1")
     app = build_app(cli_args)
 
     # Initialize the engine (this loads the model - takes time)
@@ -646,7 +649,11 @@ def _build_standalone_cli_args(argv: Optional[List[str]] = None) -> Namespace:
     ``--worker-extension-cls``, ...).
     """
     from vllm import AsyncEngineArgs as _AsyncEngineArgs
-    from vllm.entrypoints.openai.cli_args import FrontendArgs
+
+    try:  # vLLM >= 0.28.1 moved the CLI args to entrypoints.launchers
+        from vllm.entrypoints.launchers.cli_args import FrontendArgs
+    except ImportError:
+        from vllm.entrypoints.openai.cli_args import FrontendArgs
     from vllm.platforms import current_platform
     from vllm.utils.argparse_utils import FlexibleArgumentParser
 

skyrl/backends/skyrl_train/patches/megatron/patch_fa4_cute_import.py

@@ -0,0 +1,53 @@
+"""Keep a broken FlashAttention-4 (``flash_attn.cute``) install from taking down ``megatron.core``.
+
+flash-attn 2.8.x ships the FA4 ``flash_attn.cute`` package alongside FA2. It is written against
+a specific ``nvidia-cutlass-dsl`` release; vLLM >= 0.28.1 pins a newer cutlass DSL whose
+``cutlass.cute.core`` no longer has the symbols FA4 imports (``ThrMma``), so importing
+``flash_attn.cute`` raises ``AttributeError`` rather than ``ImportError``.
+
+``megatron.core.transformer.attention`` probes FA4 with ``from flash_attn.cute import
+flash_attn_varlen_func`` guarded only by ``except ImportError`` (megatron-core 0.20), so the
+``AttributeError`` escapes and every ``import megatron.core`` fails. SkyRL never uses FA4 (TE
+dispatches to FA2 or cuDNN fused attention), so when the probe import fails for any reason
+other than a plain ``ImportError`` we register ``flash_attn.cute`` as unavailable
+(``sys.modules[name] = None`` makes any later import of it raise ``ImportError``), which is
+exactly the state megatron-core handles.
+
+No-op when ``flash_attn.cute`` imports cleanly or is genuinely absent. Call it before the first
+``megatron`` import in a process; ``skyrl.backends.skyrl_train.workers.megatron`` does so on
+package import.
+
+DELETE THIS PATCH once megatron-core's FA4 probe catches ``Exception`` or the flash-attn /
+cutlass DSL pins agree again.
+"""
+
+import importlib
+import sys
+
+from loguru import logger
+
+_MODULE = "flash_attn.cute"
+_PATCHED_FLAG = "_skyrl_fa4_cute_import_patched"
+
+
+def patch_fa4_cute_import() -> bool:
+    """Return True if ``flash_attn.cute`` was marked unavailable, False if nothing was needed."""
+    if getattr(sys, _PATCHED_FLAG, False):
+        return sys.modules.get(_MODULE, True) is None
+    setattr(sys, _PATCHED_FLAG, True)
+    if _MODULE in sys.modules:
+        return sys.modules[_MODULE] is None
+    try:
+        importlib.import_module(_MODULE)
+        return False
+    except ImportError:
+        return False  # genuinely absent: megatron-core already handles this
+    except Exception as exc:  # e.g. AttributeError from a cutlass DSL mismatch
+        for name in [m for m in sys.modules if m == _MODULE or m.startswith(_MODULE + ".")]:
+            del sys.modules[name]
+        sys.modules[_MODULE] = None
+        logger.warning(
+            f"`{_MODULE}` (FlashAttention-4) failed to import with {type(exc).__name__}: {exc}. "
+            "Marking it unavailable so megatron-core falls back to FA2 / cuDNN attention."
+        )
+        return True

skyrl/backends/skyrl_train/weight_sync/sharded_rdt/rdt_vllm_register.py

@@ -22,7 +22,10 @@
 function and ImportError is swallowed on platforms without vLLM.
 
 REMOVAL: delete this module + its imports once SkyRL's pinned vLLM registers the
-``sharded_rdt`` engine in ``WeightTransferEngineFactory`` natively.
+``sharded_rdt`` engine in ``WeightTransferEngineFactory`` natively. vLLM >= 0.28.1
+(the pinned dev wheel) does register ``sharded_rdt`` natively, so on that wheel this
+module is a no-op and vLLM's own engine is used; the vendored engine below stays as the
+fallback for older wheels until the two implementations are reconciled.
 """
 
 import logging

skyrl/backends/skyrl_train/workers/megatron/__init__.py

@@ -0,0 +1,13 @@
+"""Megatron backend workers.
+
+The package import runs :func:`patch_fa4_cute_import` first: with the pinned flash-attn and
+the cutlass DSL that vLLM >= 0.28.1 pulls in, ``import megatron.core`` otherwise aborts on a
+broken FlashAttention-4 probe (see the patch module). Every Megatron entry point in SkyRL
+lives under this package, so this is the earliest common place to apply it.
+"""
+
+from skyrl.backends.skyrl_train.patches.megatron.patch_fa4_cute_import import (
+    patch_fa4_cute_import,
+)
+
+patch_fa4_cute_import()

skyrl/backends/skyrl_train/workers/megatron/glm5_next/__init__.py

@@ -0,0 +1,14 @@
+"""GLM-5.3-Flash (HF ``glm5_next``) for the Megatron backend.
+
+Shaped for an upstream home in Megatron-Bridge (``megatron/bridge/models/glm/``):
+
+- ``provider``: ``Glm5NextModelProvider`` (MLA provider + KDA / mHC / kpool fields).
+- ``layer_specs``: per-layer KDA-or-DSA, dense-or-MoE block spec built on
+  ``HyperConnectionTransformerLayer``.
+- ``dsa``: the GLM-5.3-Flash flavour of ``DSAttention`` (NoPE MLA + kpool-compressed indexer).
+- ``bridge``: ``Glm5NextBridge``, HF <-> Megatron parameter mapping for the language model of the
+  ``Glm5NextForConditionalGeneration`` checkpoint. Importing it registers the bridge.
+
+The Megatron-Core-side building blocks (mHC layer, KDA module) live in
+``skyrl.backends.skyrl_train.workers.megatron.mcore_ext``.
+"""