GLM 5.3 Flash RL Support

docker/Dockerfile

@@ -84,7 +84,14 @@ RUN pip install /tmp/wheels/flash_attn_3-*.whl
 
 RUN pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10887887bc74853f89a4de258c0702932a1c --no-deps
 
-RUN pip install flash-linear-attention==0.5.2
+# fla 0.5.2 computes BK = triton.next_power_of_2(K) inside a @triton.jit KDA kernel, which this triton
+# rejects at the first GLM-5.3-Flash trainer forward. Fixed upstream by
+# https://github.com/fla-org/flash-linear-attention/pull/1108; drop the patch once a release has it.
+COPY docker/patch/fla_kda_hoist_next_power_of_2.patch /tmp/fla_kda_hoist_next_power_of_2.patch
+RUN pip install flash-linear-attention==0.5.2 && \
+    FLA_DIR=$(python3 -c 'import importlib.util; print(importlib.util.find_spec("fla").submodule_search_locations[0])') && \
+    patch --fuzz=0 -p1 -d "$FLA_DIR" < /tmp/fla_kda_hoist_next_power_of_2.patch && \
+    rm /tmp/fla_kda_hoist_next_power_of_2.patch
 # required for DeepSeek V4
 RUN pip install tilelang==0.1.14 apache-tvm-ffi==0.1.11
 # TileKernels target-import fix: https://github.com/deepseek-ai/TileKernels/pull/25.

docker/patch/fla_kda_hoist_next_power_of_2.patch

@@ -0,0 +1,27 @@
+diff --git a/ops/kda/chunk_intra_token_parallel.py b/ops/kda/chunk_intra_token_parallel.py
+--- a/ops/kda/chunk_intra_token_parallel.py
++++ b/ops/kda/chunk_intra_token_parallel.py
+@@ -44,6 +44,7 @@
+     H: tl.constexpr,
+     HV: tl.constexpr,
+     K: tl.constexpr,
++    BK: tl.constexpr,
+     BT: tl.constexpr,
+     BC: tl.constexpr,
+     BH: tl.constexpr,
+@@ -91,7 +92,6 @@
+     Akk += bos * HV*BC
+     beta += bos * HV
+
+-    BK: tl.constexpr = triton.next_power_of_2(K)
+     o_hv = i_hg * BH + tl.arange(0, BH)
+     o_h = o_hv // G
+     o_k = tl.arange(0, BK)
+@@ -174,6 +174,7 @@
+         H=H,
+         HV=HV,
+         K=K,
++        BK=triton.next_power_of_2(K),
+         BT=BT,
+         BC=BC,
+     )

docs/models/glm/glm5-3-flash.md

@@ -28,7 +28,6 @@ GLM5.2 flagships, not a smaller cut of them.
 | Variant | `--model-name` | Layers |
 |---|---|---|
 | Full | `GLM-5.3-Flash` | 45 |
-| 8-layer slice | `GLM-5.3-Flash-8layer` | 8 |
 | 4-layer slice | `GLM-5.3-Flash-4layer` | 4 (launcher default) |
 
 ## 3. Environment Setup
@@ -83,16 +82,14 @@ python scripts/run_glm5_3_flash.py train --num-nodes 1 --num-gpus-per-node 8
 
 | Shape | TP | PP | EP | Rollout engine |
 |---|---|---|---|---|
-| 16 × 4 (full, validated) | 8 | 4 | 16 | 8 GPUs, SGLang TP 8 / EP 8 |
-| 8 × 4 (full) | 8 | 4 | 16 | 8 GPUs, SGLang TP 8 / EP 8 |
-| 6 × 4 (full) | 8 | 3 | 8 | 8 GPUs, SGLang TP 8 / EP 8 |
-| 2 × 4 or 1 × 8 (slices) | 2 | 2 | 2 | 4 GPUs, SGLang TP 4 / EP 4 |
+| 16 × 4 (full) | 8 | 4 | 16 | 8 GPUs, SGLang TP 8 / EP 8 |
+| 8 × 4 (full) | 8 | 4 | 8 | 8 GPUs, SGLang TP 8 / EP 8 |
+| 2 × 4 or 1 × 8 (4-layer slice) | 2 | 2 | 2 | 4 GPUs, SGLang TP 4 / EP 4 |
 
 The PP-4 shapes run 11 / 11 / 11 / 12 layers per stage, since 45 does not divide by 4.
 GRPO on DAPO-Math-17k, Adam at `lr 1e-6`, `max_tokens_per_gpu 8192`, full uniform recompute.
 Rollout is colocated, with the trainer offloaded to disk; both DSA paths run on tilelang and
-the KV cache is BF16. Routing replay is wired end to end, and indexer-topk replay
-(`--use-rollout-indexer-replay`) is implemented but off by default.
+the KV cache is BF16. Routing replay turns on with `--enable-r3`.
 
 ## 5. What a Healthy Run Looks Like
 

miles/backends/megatron_utils/megatron_to_hf/__init__.py

@@ -2,6 +2,7 @@
 from .deepseekv4 import convert_deepseekv4_to_hf
 from .glm4 import convert_glm4_to_hf
 from .glm4moe import convert_glm4moe_to_hf
+from .glm5_next import convert_glm5_next_to_hf
 from .inkling import convert_inkling_to_hf
 from .kimi_k3 import convert_kimi_k3_to_hf
 from .kimi_vl import convert_kimi_k25_to_hf, convert_kimivl_to_hf
@@ -34,7 +35,9 @@ def convert_to_hf(args, model_name, name, param, quantization_config=None, packe
 # TODO optimize code details
 def _convert_to_hf_core(args, model_name, name, param):
     model_name = model_name.lower()
-    if (
+    if "glm5_next" in model_name or "glm5next" in model_name:
+        converted_named_tensors = convert_glm5_next_to_hf(args, name, param)
+    elif (
         "glm4moelite" in model_name
         or "deepseekv3" in model_name
         or "glmmoedsa" in model_name

miles/backends/megatron_utils/megatron_to_hf/glm5_next.py

@@ -0,0 +1,95 @@
+"""Megatron -> HF weight-update conversion for GLM-5.3-Flash, the inverse of `miles_plugins/mbridge/glm5_next.py`.
+
+DSA indexer names are mapped here, not in `convert_deepseekv3_to_hf`, so its rope-interleave half-swap never runs
+(`qk_rope_head_dim == 0`). The three `alpha_*` of one mHC site arrive on the same rank and are buffered until they
+can be emitted as one `hc_*_scale` tensor.
+"""
+
+import re
+
+import torch
+
+from .deepseekv3 import convert_deepseekv3_to_hf
+
+_KDA_SUFFIX_MAPPING = {
+    f"self_attention.kda.{weight_name}": f"self_attn.{weight_name}"
+    for weight_name in [
+        "q_proj.weight",
+        "k_proj.weight",
+        "v_proj.weight",
+        "b_proj.weight",
+        "f_a_proj.weight",
+        "f_b_proj.weight",
+        "g_a_proj.weight",
+        "g_b_proj.weight",
+        "A_log",
+        "dt_bias",
+        "o_norm.weight",
+        "o_proj.weight",
+    ]
+}
+
+_INDEXER_SUFFIX_MAPPING = {
+    "self_attention.wq_b.weight": "self_attn.indexer.wq_b.weight",
+    "self_attention.wk.weight": "self_attn.indexer.wk.weight",
+    "self_attention.weights_proj.weight": "self_attn.indexer.weights_proj.weight",
+    "self_attention.k_norm.weight": "self_attn.indexer.k_norm.weight",
+    "self_attention.k_norm.bias": "self_attn.indexer.k_norm.bias",
+    "self_attention.index_kpool_compress_gate": "self_attn.indexer.index_kpool_compress_gate",
+    "self_attention.index_kpool_compress_ape": "self_attn.indexer.index_kpool_compress_ape",
+}
+
+_HC_SUFFIX_MAPPING = {
+    "self_attention_hyper_connection.mapping_proj.weight": "hc_attn_fn",
+    "self_attention_hyper_connection.bias": "hc_attn_base",
+    "mlp_hyper_connection.mapping_proj.weight": "hc_ffn_fn",
+    "mlp_hyper_connection.bias": "hc_ffn_base",
+}
+
+_HC_ALPHA_ORDER = ("alpha_pre", "alpha_post", "alpha_res")
+_HC_SITE_TO_SCALE = {
+    "self_attention_hyper_connection": "hc_attn_scale",
+    "mlp_hyper_connection": "hc_ffn_scale",
+}
+
+_LAYER_PATTERN = re.compile(r"module\.module\.decoder\.layers\.(\d+)\.(.+)")
+_HC_ALPHA_PATTERN = re.compile(
+    r"(self_attention_hyper_connection|mlp_hyper_connection)\.(alpha_pre|alpha_post|alpha_res)$"
+)
+
+_hc_scale_buffers: dict[tuple[str, str], dict[str, torch.Tensor]] = {}
+
+
+def _convert_hc_scale(layer_idx: str, site: str, alpha_name: str, param: torch.Tensor):
+    buffer = _hc_scale_buffers.setdefault((layer_idx, site), {})
+    buffer[alpha_name] = param
+    if len(buffer) < len(_HC_ALPHA_ORDER):
+        return []
+    _hc_scale_buffers.pop((layer_idx, site))
+    scale = torch.cat([buffer[name].reshape(1) for name in _HC_ALPHA_ORDER])
+    return [(f"model.layers.{layer_idx}.{_HC_SITE_TO_SCALE[site]}", scale)]
+
+
+def convert_glm5_next_to_hf(args, name, param):
+    match = _LAYER_PATTERN.match(name)
+    if match:
+        layer_idx, rest = match.groups()
+
+        hf_suffix = _KDA_SUFFIX_MAPPING.get(rest) or _INDEXER_SUFFIX_MAPPING.get(rest) or _HC_SUFFIX_MAPPING.get(rest)
+        if hf_suffix is not None:
+            return [(f"model.layers.{layer_idx}.{hf_suffix}", param)]
+
+        if rest == "self_attention.kda.conv1d.weight":
+            q_conv, k_conv, v_conv = param.chunk(3, dim=0)
+            return [
+                (f"model.layers.{layer_idx}.self_attn.q_conv1d.weight", q_conv),
+                (f"model.layers.{layer_idx}.self_attn.k_conv1d.weight", k_conv),
+                (f"model.layers.{layer_idx}.self_attn.v_conv1d.weight", v_conv),
+            ]
+
+        alpha_match = _HC_ALPHA_PATTERN.match(rest)
+        if alpha_match:
+            site, alpha_name = alpha_match.groups()
+            return _convert_hc_scale(layer_idx, site, alpha_name, param)
+
+    return convert_deepseekv3_to_hf(args, name, param)

miles/ray/rollout/train_data_conversion.py

@@ -136,9 +136,19 @@ def convert_samples_to_train_data(
 
     if samples[0].rollout_routed_experts is not None:
         train_data["rollout_routed_experts"] = [sample.rollout_routed_experts for sample in samples]
+    elif getattr(args, "use_rollout_routing_replay", False):
+        raise ValueError(
+            "--use-rollout-routing-replay is set but the rollout samples carry no "
+            "rollout_routed_experts: the engine response meta_info lacked 'routed_experts'."
+        )
 
     if samples[0].rollout_indexer_topk is not None:
         train_data["rollout_indexer_topk"] = [sample.rollout_indexer_topk for sample in samples]
+    elif getattr(args, "use_rollout_indexer_replay", False):
+        raise ValueError(
+            "--use-rollout-indexer-replay is set but the rollout samples carry no "
+            "rollout_indexer_topk: the engine response meta_info lacked 'indexer_topk'."
+        )
 
     if samples[0].train_metadata is not None:
         train_data["metadata"] = [sample.train_metadata for sample in samples]

miles/rollout/generate_utils/generate_endpoint_utils.py

@@ -135,7 +135,12 @@ def get_routed_experts_from_response(args, output, num_tokens: int):
     info = output["meta_info"].get("routed_experts")
     if info is None:
         return None
-    return _decode_topk_buffer(info, num_tokens, args.num_layers, -1)
+    routed_experts = _decode_topk_buffer(info, num_tokens, args.num_layers, -1)
+    assert routed_experts.size == 0 or routed_experts.any(), (
+        "routed_experts payload is all zeros: the sglang engine did not capture routed experts "
+        "(topk-bypassing --moe-runner-backend such as flashinfer_trtllm?)."
+    )
+    return routed_experts
 
 
 def get_indexer_topk_from_response(args, output, sample):
@@ -147,4 +152,9 @@ def get_indexer_topk_from_response(args, output, sample):
         "Server returned indexer_topk without indexer_topk_num_layers; "
         "sglang-miles must include the layer count in meta_info."
     )
+    expected_num_streams = getattr(args, "rollout_indexer_topk_num_streams", None)
+    assert expected_num_streams is None or num_layers == expected_num_streams, (
+        f"Server returned indexer_topk with {num_layers} streams but the model has "
+        f"{expected_num_streams} indexer layers; replaying it would map streams to the wrong layers."
+    )
     return _decode_topk_buffer(info, len(sample.tokens) - 1, num_layers, -1)

miles/rollout/sglang_rollout.py

@@ -268,6 +268,10 @@ async def generate(
             f"prompt_tokens={output['meta_info'].get('prompt_tokens')} response={len(new_response_tokens)} "
             f"unexpanded_tokens={len(sample.tokens)}"
         )
+        assert _re.size == 0 or _re.any(), (
+            "routed_experts payload is all zeros: the sglang engine did not capture routed experts "
+            "(topk-bypassing --moe-runner-backend such as flashinfer_trtllm?)."
+        )
         sample.rollout_routed_experts = _re.reshape(_ntok, args.num_layers, _topk)
     if "indexer_topk" in output["meta_info"]:
         sample.rollout_indexer_topk = get_indexer_topk_from_response(args, output, sample)

miles/utils/arguments.py

@@ -2600,14 +2600,20 @@ def parse_args(add_custom_arguments=None, entry="train", preprocess_args=None):
 
         args = megatron_parse_args(extra_args_provider=add_miles_arguments)
         args.compress_ratios = None
+        args.rollout_indexer_topk_num_streams = None
         if args.hf_checkpoint:
             hf_config = load_hf_config(args.hf_checkpoint)
             args.compress_ratios = getattr(hf_config, "compress_ratios", None)
             hf_validate_args(args, hf_config)
 
             if is_dsa(hf_config):
-                args.indexer_rope_interleave = bool(getattr(hf_config, "indexer_rope_interleave", False))
+                getter = getattr(hf_config, "get_text_config", None)
+                text_config = (getter() if callable(getter) else getattr(hf_config, "text_config", None)) or hf_config
+                args.indexer_rope_interleave = bool(getattr(text_config, "indexer_rope_interleave", False))
                 logger.info(f"Setting indexer_rope_interleave: {args.indexer_rope_interleave} into args")
+                linear_attn_config = getattr(text_config, "linear_attn_config", None)
+                kda_layers = set((linear_attn_config or {}).get("kda_layers") or [])
+                args.rollout_indexer_topk_num_streams = text_config.num_hidden_layers - len(kda_layers)
 
         # TODO: unify this .rank and .world_size w/ indep_dp logics
         args.rank = 0

miles/utils/hf_utils/config.py

@@ -62,6 +62,13 @@ class _HFConfigAlias:
         compat_class_name="Qwen4ExpConfig",
         auto_model_classes=(),
     ),
+    _HFConfigAlias(
+        model_type="glm5_next",
+        base_module="transformers.models.glm4v_moe.configuration_glm4v_moe",
+        base_class="Glm4vMoeConfig",
+        compat_class_name="Glm5NextConfig",
+        auto_model_classes=(),
+    ),
 )
 
 _REGISTERED_ALIASES: set[str] = set()
@@ -128,7 +135,7 @@ def load_hf_config(
 
 
 def is_dsa(hf_config) -> bool:
-    return getattr(hf_config, "model_type", None) in ("deepseek_v32", "glm_moe_dsa")
+    return getattr(hf_config, "model_type", None) in ("deepseek_v32", "glm_moe_dsa", "glm5_next")
 
 
 # Written by HF exports after all ranks finish, so consumers can tell finished from partial.

miles_plugins/mbridge/__init__.py

@@ -3,6 +3,7 @@
 from .glm4 import GLM4Bridge
 from .glm4moe import GLM4MoEBridge
 from .glm4moe_lite import GLM4MoELiteBridge
+from .glm5_next import Glm5NextBridge
 from .inkling import InklingBridge
 from .joyai_llm_flash import JoyAILLMFlashBridge
 from .kimi_k3 import KimiK3Bridge
@@ -20,6 +21,7 @@
     "Qwen38NextBridge",
     "MimoBridge",
     "DeepseekV32Bridge",
+    "Glm5NextBridge",
     "DeepseekV4Bridge",
     "JoyAILLMFlashBridge",
     "InklingBridge",