[model, peft] feat: GLM-5 DSA option: tilelang (2/2)

src/megatron/bridge/models/__init__.py

@@ -59,7 +59,7 @@
 from megatron.bridge.models.glm import (
     GLM45Bridge,
 )
-from megatron.bridge.models.glm_moe_dsa import (
+from megatron.bridge.models.glm5 import (
     GLM5Bridge,
 )
 from megatron.bridge.models.glm_vl import (

src/megatron/bridge/models/glm5/__init__.py

@@ -0,0 +1,51 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""GLM-5.x (``glm_moe_dsa``: MoE + MLA + DeepSeek Sparse Attention) Megatron-Bridge model.
+
+Two DSA sparse-MLA kernel backends are selectable via the provider attribute
+``config.dsa_attention_backend`` (set from the miles ``--dsa-attention-backend`` arg under
+``--megatron-to-hf-mode bridge``). The choice is **orthogonal to the model version and to LoRA** --
+BOTH backends support GLM-5.1 *and* GLM-5.2 (DSA cross-layer index sharing), full or LoRA:
+
+  * ``"megatron"`` (default) -- the portable *unfused* megatron-core DSA kernels
+    (``DSAttention`` / ``CrossLayerDSAttention`` in ``cross_layer_dsa_dispatch.py``). No extra dependencies.
+    Works with both the ``bshd`` and ``thd`` query layouts; ``thd`` is the preferred,
+    activation-recompute-safe carrier, while ``bshd`` + activation recompute is rejected at forward
+    time (the ``cross_layer_dsa_dispatch.py`` forward guard).
+
+  * ``"tilelang"`` -- the vendored *fused* TileLang kernels (``SparseMLA`` + ``lighting_indexer`` under
+    ``tilelang/``, driven by ``TileLangMLASelfAttention`` in ``tilelang_mla.py``). Matches slime's rollout
+    kernels for rollout<->train numerical parity, including R3 indexer replay. Requires the optional
+    ``tilelang`` dependency (imported lazily, so the default path stays dependency-free) and the
+    ``thd`` (packed) layout (``--qkv-format thd``). **Training/forward-only**: it asserts
+    ``inference_context is None`` (no KV cache) and cannot serve generation -- the rollout is always
+    served by sglang; tilelang only matches sglang numerically on the *train* side.
+
+Both backends are LoRA-capable. The DSA indexer (``wq_b`` / ``wk`` / ``weights_proj``) is excluded
+from LoRA by default in the miles launcher. On the tilelang backend the indexer adapters get no
+gradient (the fused ``lighting_indexer`` returns only discrete top-k and computes no indexer loss),
+so training it there is a genuine no-op; on the default/unfused backend the indexer adapter *would*
+get a tiny aux-loss gradient (~1e-5, ``dsa_indexer_loss_coeff=0.001``), so excluding it there is a
+deliberate choice. GLM-5.2's cross-layer index sharing (``index_topk_freq`` /
+``index_skip_topk_offset``) is read from the HF config by ``GLM5Bridge`` and honored by both
+backends -- no extra CLI args.
+"""
+
+from megatron.bridge.models.glm5.glm5_bridge import GLM5Bridge
+
+
+__all__ = [
+    "GLM5Bridge",
+]

src/megatron/bridge/models/glm_moe_dsa/glm5_provider.py → src/megatron/bridge/models/glm5/glm5_provider.py

@@ -12,9 +12,21 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-"""GLM5 uses MLAModelProvider directly. This module is kept for import compatibility."""
+"""GLM5 model provider: MLAModelProvider plus the GLM-5 DSA configuration fields."""
 
-from megatron.bridge.models.mla_provider import MLAModelProvider as GLM5ModelProvider  # noqa: F401
+from dataclasses import dataclass
+
+from megatron.bridge.models.mla_provider import MLAModelProvider
+
+
+@dataclass
+class GLM5ModelProvider(MLAModelProvider):
+    """GLM-5 (glm_moe_dsa) provider: MLA plus DeepSeek Sparse Attention."""
+
+    # DSA sparse-MLA kernel backend; set from the miles --dsa-attention-backend arg. Declared
+    # as a real field (rather than an ad-hoc attribute) so it survives any fields-based config
+    # copy and reaches every module's config without caller-side propagation.
+    dsa_attention_backend: str = "megatron"
 
 
 __all__ = ["GLM5ModelProvider"]

src/megatron/bridge/models/glm5/megatron/__init__.py

@@ -0,0 +1,30 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License").
+
+"""GLM-5 UNFUSED DSA backend (goes through megatron-core / megatron-bridge).
+
+Selected when ``dsa_attention_backend != "tilelang"`` (the default). The unfused
+sparse-MLA path reuses megatron-core's experimental ``DSAttention`` (lightning
+indexer + ``unfused_dsa_fn``); GLM-5.2 cross-layer index-sharing is layered on
+top by ``CrossLayerDSAttention`` in ``../cross_layer_dsa_dispatch.py``.
+
+Centralising the megatron-core DSA imports here marks the unfused backend's
+single entry point (mirrors how ``../tilelang/`` is the tilelang entry point).
+"""
+
+from megatron.core.transformer.experimental_attention_variant.dsa import (
+    DSAIndexerLossAutoScaler,
+    DSAIndexerLossLoggingHelper,
+    DSAttention,
+    FusedDSAIndexerLoss,
+    unfused_dsa_fn,
+)
+
+__all__ = [
+    "DSAttention",
+    "FusedDSAIndexerLoss",
+    "unfused_dsa_fn",
+    "DSAIndexerLossAutoScaler",
+    "DSAIndexerLossLoggingHelper",
+]

src/megatron/bridge/models/glm5/tilelang/__init__.py

@@ -0,0 +1,27 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Fused TileLang DSA kernels for GLM ``glm_moe_dsa`` (vendored from THUDM/slime).
+
+Provides the ``tilelang`` sparse-attention backend: ``SparseMLA`` (sparse-MLA attention) and
+``lighting_indexer`` (the DSA indexer), both with fwd+bwd TileLang kernels. Importing this
+package pulls in the optional ``tilelang`` dependency, so
+``cross_layer_dsa_dispatch.py`` imports it lazily — only when ``config.dsa_attention_backend == "tilelang"``
+— keeping the default unfused (``megatron``) path dependency-free.
+"""
+from .indexer import generate_varlen_mask_params, lighting_indexer
+from .sparse_mla import SparseMLA
+
+
+__all__ = ["SparseMLA", "generate_varlen_mask_params", "lighting_indexer"]

src/megatron/bridge/models/glm5/tilelang/indexer.py

@@ -0,0 +1,111 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# ruff: noqa: D101, D103, E741, F401
+#
+# Vendored from THUDM/slime (miles_plugins/models/glm5/ops) for the GLM DSA fused attention
+# backend. Kept close to upstream; the only change is removing the miles-only indexer replay
+# hook. Imported lazily by glm5/tilelang/tilelang_mla.py only when dsa_attention_backend="tilelang".
+
+import torch
+
+from .tilelang_indexer_bwd import indexer_bwd_interface
+from .tilelang_indexer_fwd import indexer_fwd_interface
+
+
+def pytorch_extract_topk_scores(logits, topk_indices, dim=-1):
+    valid_mask = topk_indices != -1
+    safe_indices = topk_indices.clamp(min=0).to(torch.int64)
+    scores = torch.gather(logits, dim=dim, index=safe_indices)
+    scores = torch.where(valid_mask, scores, float("-inf"))
+    return scores
+
+
+def _original_topk(logits, topk):
+    # Short sequence (seq_len_kv < index_topk): the indexer degenerates to dense. torch.topk cannot
+    # select more entries than exist ("selected index k out of range"), so cap k and pad the
+    # selection back out to the fixed `topk` width with -1 (invalid). This matches the
+    # rollout-captured indexer top-k shape that R3 replay asserts ([n_tokens, topk]) and that
+    # SparseMLA expects; -1 is the same sentinel masked_fill uses for out-of-window picks, which
+    # downstream ignores. The long-sequence path (k == topk) is unchanged.
+    k = min(topk, logits.shape[-1])
+    score, indices = torch.topk(logits, k, dim=-1)
+    indices = indices.to(torch.int32).masked_fill(score == -torch.inf, -1)
+    if k < topk:
+        pad = torch.full((*indices.shape[:-1], topk - k), -1, dtype=torch.int32, device=indices.device)
+        indices = torch.cat([indices, pad], dim=-1)
+    return indices
+
+
+class IndexerFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(
+        ctx,
+        index_q: torch.Tensor,
+        index_k: torch.Tensor,
+        weights: torch.Tensor,
+        cu_seqlen_ks: torch.Tensor,
+        cu_seqlen_ke: torch.Tensor,
+        logits: torch.Tensor,
+        topk_indices: torch.Tensor,
+    ):
+        index_score = pytorch_extract_topk_scores(logits, topk_indices)
+        ctx.save_for_backward(index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices)
+        return index_score
+
+    @staticmethod
+    def backward(ctx, grad_scores):
+        index_q, index_k, weights, cu_seqlen_ks, cu_seqlen_ke, topk_indices = ctx.saved_tensors
+        grad_q, grad_w, grad_k = indexer_bwd_interface(index_q, weights, index_k, topk_indices, grad_scores)
+        return grad_q, grad_k, grad_w, None, None, None, None
+
+
+def lighting_indexer(
+    index_q: torch.Tensor,
+    index_k: torch.Tensor,
+    weights: torch.Tensor,
+    cu_seqlen_ks: torch.Tensor,
+    cu_seqlen_ke: torch.Tensor,
+    topk: int,
+    topk_indices: torch.Tensor | None = None,
+):
+    weights_2d = weights.squeeze(-1)
+    logits = indexer_fwd_interface(index_q, index_k, weights_2d, cu_seqlen_ks, cu_seqlen_ke, clean_logits=True)
+
+    if topk_indices is None:
+        # R3 indexer replay (matched DSA top-k between rollout & training, arxiv 2510.11370): route
+        # the selection through miles' indexer_replay_manager when it is present + enabled, mirroring
+        # slime's indexer.py so the fused backend records/replays the indexer top-k instead of
+        # recomputing it. Guarded so the vendored kernel still imports + runs standalone (no miles),
+        # in which case it falls back to a plain top-k -- identical to the previous behaviour.
+        try:
+            from miles.utils.replay_base import indexer_replay_manager
+
+            topk_fn = indexer_replay_manager.get_topk_fn(_original_topk, return_probs=False)
+        except ImportError:
+            topk_fn = _original_topk
+        topk_indices = topk_fn(logits, topk)
+
+    index_score = IndexerFunction.apply(index_q, index_k, weights_2d, cu_seqlen_ks, cu_seqlen_ke, logits, topk_indices)
+    return index_score, topk_indices
+
+
+def generate_varlen_mask_params(cu_seqlens):
+    seq_len = cu_seqlens[-1].item()
+    q_indices = torch.arange(0, seq_len, device=cu_seqlens.device)
+    seq_indices = torch.searchsorted(cu_seqlens, q_indices, right=True) - 1
+    starts = cu_seqlens[seq_indices]
+    ends = q_indices + 1
+    assert torch.all((ends - starts) > 0)
+    return starts, ends

src/megatron/bridge/models/glm5/tilelang/sparse_mla.py

@@ -0,0 +1,64 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# ruff: noqa: D101, D103, E741, F401
+#
+# Vendored from THUDM/slime (miles_plugins/models/glm5/ops) for the GLM DSA fused attention
+# backend. Kept close to upstream; the only change is removing the miles-only indexer replay
+# hook. Imported lazily by glm5/tilelang/tilelang_mla.py only when dsa_attention_backend="tilelang".
+
+import torch
+
+from .tilelang_sparse_mla_bwd import sparse_mla_bwd
+from .tilelang_sparse_mla_fwd import sparse_mla_fwd_interface
+
+
+class SparseMLA(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, q, kv, indices, scaling):
+        """
+        Args:
+            q: Query tensor (seq_len, heads, dim_plus_tail_dim)
+            kv: Key-Value tensor (seq_len_kv, kv_group, dim_plus_tail_dim)
+            indices: Sparse indices tensor (seq_len, kv_group, topk)
+
+        Returns:
+            out: Output tensor (seq_len, heads, dim)
+        """
+        indices = indices.contiguous()
+        q, kv = q.contiguous(), kv.contiguous()
+        ctx.scaling = scaling
+        tl_out, tl_lse = sparse_mla_fwd_interface(q, kv, indices, sm_scale=scaling)
+
+        # Save tensors for backward pass
+        ctx.save_for_backward(q, kv, indices, tl_out, tl_lse)
+
+        return tl_out, tl_lse
+
+    @staticmethod
+    def backward(ctx, grad_output, grad_lse):
+        """
+        Args:
+            grad_output: Gradient of the loss with respect to output
+
+        Returns:
+            Gradients for q, kv, and indices (None for indices)
+        """
+        q, kv, indices, tl_out, tl_lse = ctx.saved_tensors
+        scaling = ctx.scaling
+
+        tl_dq, tl_dkv = sparse_mla_bwd(q, kv, tl_out, grad_output.contiguous(), indices, tl_lse, sm_scale=scaling)
+
+        # Return gradients for each input (None for indices as it's not differentiable)
+        return tl_dq, tl_dkv, None, None