cp: feat: cherry-pick Ling-V3 Tiny support to main

megatron/core/inference/contexts/dynamic_context.py

@@ -453,18 +453,19 @@ def __init__(self, model_config: TransformerConfig, inference_config: InferenceC
                     "boundaries are not rounded between decode chunks."
                 )
 
-            # Mamba and GDN use the same slot-indexed recurrent-state cache contract. Build
-            # one map in global layer order; independently generated per-symbol maps both
-            # start at zero and would alias if they were simply unioned.
-            attention_layer_map, dsa_layer_map = operator.itemgetter(
-                Symbols.ATTENTION, Symbols.DS_ATTENTION
-            )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list))
-            recurrent_layer_map = {}
-            for global_layer_idx, layer_type in enumerate(
-                mamba_inference_state_config.layer_type_list
-            ):
-                if layer_type in (Symbols.MAMBA, Symbols.GDN):
-                    recurrent_layer_map[global_layer_idx] = len(recurrent_layer_map)
+            # For hybrid models, the layer map converts the global layer index to the
+            # corresponding attention layer index or Mamba layer index depending on the
+            # layer type.
+            attention_layer_map, dsa_layer_map, gdn_layer_map, kda_layer_map, mamba_layer_map = (
+                operator.itemgetter(
+                    Symbols.ATTENTION, Symbols.DS_ATTENTION, Symbols.GDN, Symbols.KDA, Symbols.MAMBA
+                )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list))
+            )
+
+            if len(gdn_layer_map) > 0:
+                raise NotImplementedError("GDN layers are not supported for inference.")
+            if len(kda_layer_map) > 0:
+                raise NotImplementedError("KDA layers are not supported for inference.")
 
             self.num_attention_layers = len(attention_layer_map) + len(dsa_layer_map)
             self.num_mamba_layers = len(recurrent_layer_map)
@@ -3188,7 +3189,7 @@ def check_availability(self, req: DynamicInferenceRequest) -> Tuple[bool, bool,
         if self.is_hybrid_model and self.kv_block_allocator.enable_handoff_pinning:
             request_can_be_added &= self.mamba_metadata.mamba_state_free_slot_count > 0
 
-        matched_block_ids, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length = (
+        _, num_blocks_from_pool, _, _, _, effective_prefill_chunk_length = (
             self._compute_prefix_match(req, req.remaining_prompt_length)
         )
 

megatron/core/models/gpt/experimental_attention_variant_module_specs.py

@@ -5,7 +5,12 @@
 
 from megatron.core.fusions.fused_bias_dropout import get_bias_dropout_add
 from megatron.core.models.backends import BackendSpecProvider
-from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNet2, GatedDeltaNetSubmodules
+from megatron.core.ssm.gated_delta_net import (
+    GatedDeltaNet,
+    GatedDeltaNetSubmodules,
+    KimiDeltaAttention,
+    KimiDeltaAttentionSubmodules,
+)
 from megatron.core.transformer.enums import AttnMaskType, LayerType
 from megatron.core.transformer.experimental_attention_variant.absorbed_mla import (
     AbsorbedMLASelfAttention,
@@ -58,7 +63,7 @@
 ##########
 
 # Canonical ``experimental_attention_variant`` names served by the gated delta net family.
-GDN_ATTENTION_VARIANTS = ("gdn", "gdn2")
+GDN_ATTENTION_VARIANTS = ("gdn", "kda")
 
 # Deprecated ``experimental_attention_variant`` spellings mapped to their canonical name.
 _DEPRECATED_ATTENTION_VARIANT_ALIASES = {"gated_delta_net": "gdn"}
@@ -72,25 +77,33 @@
 def get_gated_delta_net_module_spec(
     config: TransformerConfig, backend: BackendSpecProvider = None
 ) -> ModuleSpec:
-    """Build module spec for GatedDeltaNet attention."""
+    """Build a module spec for a GDN-family attention variant."""
 
     if backend is None:
         backend = _get_backend_spec_provider(config=config)
 
     rms_norm = config.normalization == "RMSNorm"
-    # gdn2 reuses the GDN submodules and spec structure with the GatedDeltaNet2 module.
-    gdn_module = (
-        GatedDeltaNet2 if config.experimental_attention_variant == "gdn2" else GatedDeltaNet
-    )
-    attention = ModuleSpec(
-        module=gdn_module,
-        submodules=GatedDeltaNetSubmodules(
-            in_proj=backend.column_parallel_layer_norm_linear(),
-            out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False),
-            out_proj=backend.row_parallel_linear(),
-        ),
-        metainfo={"fuse_input_layernorm": True},
-    )
+    if config.experimental_attention_variant == "kda":
+        attention = ModuleSpec(
+            module=KimiDeltaAttention,
+            submodules=KimiDeltaAttentionSubmodules(
+                in_proj=backend.column_parallel_linear(),
+                beta_proj=backend.column_parallel_linear(),
+                out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False),
+                out_proj=backend.row_parallel_linear(),
+            ),
+            metainfo={"fuse_input_layernorm": False},
+        )
+    else:
+        attention = ModuleSpec(
+            module=GatedDeltaNet,
+            submodules=GatedDeltaNetSubmodules(
+                in_proj=backend.column_parallel_layer_norm_linear(),
+                out_norm=backend.layer_norm(rms_norm=rms_norm, for_qk=False),
+                out_proj=backend.row_parallel_linear(),
+            ),
+            metainfo={"fuse_input_layernorm": True},
+        )
     return attention
 
 
@@ -349,19 +362,7 @@ def get_transformer_block_with_experimental_attention_variant_spec(
 def normalize_experimental_attention_variant(
     experimental_attention_variant: Optional[str],
 ) -> Optional[str]:
-    """Resolve a deprecated ``experimental_attention_variant`` spelling to its canonical name.
-
-    ``gated_delta_net`` is the deprecated spelling of ``gdn``. Passing it emits a
-    ``DeprecationWarning`` and returns the canonical name so that every downstream
-    consumer only has to handle ``gdn``.
-
-    Args:
-        experimental_attention_variant: The configured variant name, possibly a
-            deprecated alias.
-
-    Returns:
-        The canonical variant name, or the argument unchanged when it is not an alias.
-    """
+    """Resolve a deprecated attention-variant spelling to its canonical name."""
     canonical = _DEPRECATED_ATTENTION_VARIANT_ALIASES.get(experimental_attention_variant)
     if canonical is None:
         return experimental_attention_variant
@@ -376,11 +377,7 @@ def normalize_experimental_attention_variant(
 
 
 def is_gated_delta_net_variant(experimental_attention_variant: Optional[str]) -> bool:
-    """Check if the experimental attention variant is served by a gated delta net layer.
-
-    Accepts the deprecated ``gated_delta_net`` spelling without warning; use
-    :func:`normalize_experimental_attention_variant` to emit the deprecation notice.
-    """
+    """Return whether a name selects a GDN-family attention implementation."""
     canonical = _DEPRECATED_ATTENTION_VARIANT_ALIASES.get(
         experimental_attention_variant, experimental_attention_variant
     )

megatron/core/models/gpt/gpt_layer_specs.py

@@ -125,6 +125,7 @@ def get_gpt_layer_with_inference_submodules(
                     linear_kv_down_proj=backend.linear(),
                     linear_kv_up_proj=linear_kv_up_proj,
                     core_attention=backend.core_attention(),
+                    linear_gate=backend.column_parallel_linear(),
                     linear_proj=backend.row_parallel_linear(),
                     q_layernorm=IdentityOp,
                     kv_layernorm=IdentityOp,
@@ -304,6 +305,7 @@ def get_gpt_layer_with_transformer_engine_submodules(
                     linear_kv_down_proj=backend.linear(),
                     linear_kv_up_proj=linear_kv_up_proj,
                     core_attention=backend.core_attention(),
+                    linear_gate=backend.column_parallel_linear(),
                     linear_proj=backend.row_parallel_linear(),
                     q_layernorm=IdentityOp,
                     kv_layernorm=IdentityOp,
@@ -424,6 +426,7 @@ def get_gpt_layer_local_submodules(
                     linear_kv_down_proj=backend.column_parallel_linear(),
                     linear_kv_up_proj=backend.column_parallel_linear(),
                     core_attention=backend.core_attention(),
+                    linear_gate=backend.column_parallel_linear(),
                     linear_proj=backend.row_parallel_linear(),
                     q_layernorm=qk_norm if qk_layernorm else IdentityOp,
                     kv_layernorm=qk_norm if qk_layernorm else IdentityOp,

megatron/core/models/hybrid/hybrid_block.py

@@ -45,9 +45,13 @@ class HybridStackSubmodules:
 
     mamba_layer: Union[ModuleSpec, type] = IdentityOp
     gdn_layer: Union[ModuleSpec, type] = IdentityOp
+    kda_layer: Union[ModuleSpec, type] = IdentityOp
     attention_layer: Union[ModuleSpec, type] = IdentityOp
     dsa_layer: Union[ModuleSpec, type] = IdentityOp
     mla_layer: Union[ModuleSpec, type] = IdentityOp
+    csa_layer: Union[ModuleSpec, type] = IdentityOp
+    hca_layer: Union[ModuleSpec, type] = IdentityOp
+    window_layer: Union[ModuleSpec, type] = IdentityOp
     mlp_layer: Union[ModuleSpec, type] = IdentityOp
     moe_layer: Union[ModuleSpec, type] = IdentityOp
     mtp_block_spec: Optional[ModuleSpec] = None
@@ -183,6 +187,29 @@ def __init__(
                         name=(name + f".layers.{i}") if name is not None else None,
                     )
                 elif layer_type == LayerSymbols.MLA:
+                    layer = build_module(
+                        submodules.mla_layer,
+                        config=self.config,
+                        layer_number=layer_number,
+                        pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
+                        add_layer_offset=False,
+                        pp_layer_offset=pp_layer_offset,
+                        name=(name + f".layers.{i}") if name is not None else None,
+                    )
+                elif layer_type == LayerSymbols.CSA:
+                    # DSv4 Compressed Sparse Attention (compress_ratio fixed by the spec).
+                    layer = build_module(
+                        submodules.csa_layer,
+                        config=self.config,
+                        layer_number=layer_number,
+                        pg_collection=pg_collection,
+                        is_mtp_layer=is_mtp_layer,
+                        add_layer_offset=False,
+                        pp_layer_offset=pp_layer_offset,
+                    )
+                elif layer_type == LayerSymbols.HCA:
+                    # DSv4 Heavily Compressed Attention (compress_ratio fixed by the spec).
                     layer = build_module(
                         submodules.mla_layer,
                         config=self.config,
@@ -230,6 +257,15 @@ def __init__(
                         pp_layer_offset=pp_layer_offset,
                         name=(name + f".layers.{i}") if name is not None else None,
                     )
+                elif layer_type == LayerSymbols.KDA:
+                    layer = build_module(
+                        submodules.kda_layer,
+                        config=self.config,
+                        layer_number=layer_number,
+                        pg_collection=pg_collection,
+                        add_layer_offset=False,
+                        name=(name + f".layers.{i}") if name is not None else None,
+                    )
                 else:
                     raise ValueError("unexpected layer_type")
             self.layers.append(layer)

megatron/core/models/hybrid/hybrid_layer_allocation.py

@@ -16,15 +16,21 @@ class Symbols:
 
     MAMBA = "M"
     GDN = 'G'
+    KDA = 'K'
     ATTENTION = "*"
     DS_ATTENTION = "D"
     MLA = "+"
+    CSA = "C"  # DSv4 Compressed Sparse Attention (compress_ratio=4)
+    HCA = "H"  # DSv4 Heavily Compressed Attention (compress_ratio=128)
+    WINDOW = "W"  # DSv4 sliding-window-only attention (compress_ratio=0; no compressor/indexer)
     MLP = "-"
     MOE = 'E'
     PIPE = '|'
     MTP_SEPARATOR = "/"
-    VALID_LAYERS = {MAMBA, GDN, ATTENTION, DS_ATTENTION, MLA, MLP, MOE}
-    ATTENTION_LAYERS = {ATTENTION, DS_ATTENTION, MLA}
+    VALID_LAYERS = {MAMBA, GDN, KDA, ATTENTION, DS_ATTENTION, MLA, CSA, HCA, WINDOW, MLP, MOE}
+    ATTENTION_LAYERS = {ATTENTION, DS_ATTENTION, MLA, CSA, HCA, WINDOW}
+    # MLA-based attention layers (incompatible with standard '*' attention in one model).
+    MLA_ATTENTION = {MLA, DS_ATTENTION, CSA, HCA, WINDOW}
 
     @classmethod
     def name_sorted_valid_layer_symbols(cls) -> list[str]:
@@ -175,10 +181,10 @@ def get_hybrid_layer_counts(pattern: str) -> Dict[str, int]:
 
     Examples:
         >>> get_hybrid_layer_counts("M*M*")
-        {'*': 2, 'G': 0, 'D': 0, 'M': 2, '-': 0, 'E': 0}
+        {'*': 2, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'K': 0, 'M': 2, '+': 0, '-': 0, 'E': 0, 'W': 0}
 
         >>> get_hybrid_layer_counts("M-M-|M-M*-/MM/MM")
-        {'*': 1, 'G': 0, 'D': 0, 'M': 8, '-': 4, 'E': 0}
+        {'*': 1, 'C': 0, 'D': 0, 'G': 0, 'H': 0, 'K': 0, 'M': 8, '+': 0, '-': 4, 'E': 0, 'W': 0}
     """
     parsed = parse_hybrid_pattern(pattern)
     counts = {symbol: 0 for symbol in Symbols.name_sorted_valid_layer_symbols()}
@@ -294,9 +300,12 @@ def _validate_pattern(pattern: str, pattern_name: str, allow_pipe: bool = False)
                 f"Valid symbols are: {valid_chars}"
             )
 
-    # Disallow Attention + MLA/DSA hybridity.
-    if Symbols.ATTENTION in pattern and (Symbols.DS_ATTENTION in pattern or Symbols.MLA in pattern):
-        raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
+    # Disallow standard Attention ('*') mixed with any MLA-based attention (+/D/C/H/W).
+    # MLA variants (DSA / CSA / HCA / Window) may freely coexist with each other.
+    if Symbols.ATTENTION in pattern and any(s in pattern for s in Symbols.MLA_ATTENTION):
+        raise ValueError(
+            "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model"
+        )
 
 
 def validate_segment_layers(segment: str) -> List[str]:
@@ -322,9 +331,11 @@ def validate_segment_layers(segment: str) -> List[str]:
                 f"one of {Symbols.VALID_LAYERS}"
             )
 
-    # Disallow Attention + MLA/DSA hybridity.
-    if Symbols.ATTENTION in segment and (Symbols.DS_ATTENTION in segment or Symbols.MLA in segment):
-        raise ValueError("Not supported to have both Attention and MLA/DSA in one model")
+    # Disallow standard Attention ('*') mixed with any MLA-based attention (+/D/C/H/W).
+    if Symbols.ATTENTION in segment and any(s in segment for s in Symbols.MLA_ATTENTION):
+        raise ValueError(
+            "Not supported to have both Attention and MLA/DSA/CSA/HCA/Window in one model"
+        )
 
     return layer_type_list
 

megatron/core/models/hybrid/hybrid_layer_specs.py

@@ -15,7 +15,12 @@
     get_moe_module_spec,
 )
 from megatron.core.models.hybrid.hybrid_block import HybridStack, HybridStackSubmodules
-from megatron.core.ssm.gated_delta_net import GatedDeltaNet, GatedDeltaNetSubmodules
+from megatron.core.ssm.gated_delta_net import (
+    GatedDeltaNet,
+    GatedDeltaNetSubmodules,
+    KimiDeltaAttention,
+    KimiDeltaAttentionSubmodules,
+)
 from megatron.core.ssm.gated_delta_product import (
     GatedDeltaProductMixer,
     GatedDeltaProductMixerSubmodules,
@@ -133,6 +138,22 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj):
                 self_attn_bda=get_bias_dropout_add,
             ),
         ),
+        kda_layer=ModuleSpec(
+            module=TransformerLayer,
+            submodules=TransformerLayerSubmodules(
+                input_layernorm=TENorm,
+                self_attention=ModuleSpec(
+                    module=KimiDeltaAttention,
+                    submodules=KimiDeltaAttentionSubmodules(
+                        in_proj=TEColumnParallelLinear,
+                        beta_proj=TEColumnParallelLinear,
+                        out_norm=TENorm,
+                        out_proj=TERowParallelLinear,
+                    ),
+                ),
+                self_attn_bda=get_bias_dropout_add,
+            ),
+        ),
         # Started with spec from gpt_layer_specs.py (with MLP removed)
         # Using the TE spec because we had problems getting the non-TE spec
         # working
@@ -200,6 +221,7 @@ def _get_gated_delta_product_mamba_layer_spec(in_proj, out_proj):
                         linear_kv_down_proj=TELinear,
                         linear_kv_up_proj=TEColumnParallelLinear,
                         core_attention=TEDotProductAttention,
+                        linear_gate=TEColumnParallelLinear,
                         linear_proj=TERowParallelLinear,
                         q_layernorm=IdentityOp,
                         kv_layernorm=IdentityOp,