[perf] feat: add GDN (Gated DeltaNet) FLOPs calculator

Port the GDN FLOPs formula from Megatron-LM (PR #1989) into Bridge's
flop_utils.py so that Qwen3.5 VL and Qwen3-Next models report accurate
throughput numbers instead of treating all layers as standard attention.

When experimental_attention_variant="gated_delta_net" is set on the model
config, transformer_flops() now:
- Parses linear_attention_freq to split layers into GDN vs standard attention
- Computes per-layer GDN cost (in_proj + conv1d + gated delta rule + out_proj)
- Produces a weighted self_attn_term over both layer types

Signed-off-by: Chen Cui <chcui@nvidia.com>

Author: cuichenx

src/megatron/bridge/training/utils/flop_utils.py

@@ -382,6 +382,52 @@ def transformer_flops():
                 )
             )
 
+        # Handle GDN (Gated DeltaNet) hybrid attention variant.
+        # When experimental_attention_variant is "gated_delta_net", a fraction of the
+        # layers use GDN instead of standard attention. Override self_attn_term with a
+        # weighted sum of GDN and standard-attention per-layer costs.
+        experimental_attention_variant = getattr(cfg.model, "experimental_attention_variant", None)
+        if experimental_attention_variant == "gated_delta_net":
+            linear_attention_freq = getattr(cfg.model, "linear_attention_freq", None)
+            if linear_attention_freq is not None:
+                if isinstance(linear_attention_freq, int):
+                    linear_attention_pattern = [
+                        0 if ((i + 1) % linear_attention_freq == 0) else 1 for i in range(num_layers)
+                    ]
+                elif isinstance(linear_attention_freq, list):
+                    linear_attention_pattern = list(linear_attention_freq)
+                else:
+                    linear_attention_pattern = [1] * num_layers
+
+                num_gdn_layers = sum(linear_attention_pattern)
+                num_standard_attn_layers = num_layers - num_gdn_layers
+
+                standard_self_attn_per_layer = self_attn_term / num_layers if num_layers > 0 else 0
+
+                qk_head_dim = getattr(cfg.model, "linear_key_head_dim", 128)
+                v_head_dim = getattr(cfg.model, "linear_value_head_dim", 128)
+                num_qk_heads = getattr(cfg.model, "linear_num_key_heads", 16)
+                num_v_heads = getattr(cfg.model, "linear_num_value_heads", 48)
+                conv_kernel_dim = getattr(cfg.model, "linear_conv_kernel_dim", 4)
+
+                qk_dim = qk_head_dim * num_qk_heads
+                v_dim = v_head_dim * num_v_heads
+
+                gdn_self_attn_per_layer = (
+                    3
+                    * 2
+                    * (
+                        cfg.model.hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads)
+                        + conv_kernel_dim * (2 * qk_dim + v_dim)
+                        + num_v_heads * (v_head_dim**2) * 4
+                        + cfg.model.hidden_size * v_dim
+                    )
+                )
+
+                self_attn_term = (
+                    gdn_self_attn_per_layer * num_gdn_layers + standard_self_attn_per_layer * num_standard_attn_layers
+                )
+
         padded_vocab_size = calculate_padded_vocab_size(
             cfg.model.vocab_size,
             cfg.model.make_vocab_size_divisible_by,

tests/unit_tests/training/utils/test_flop_utils.py

@@ -61,6 +61,14 @@ class MockModelConfig:
     group_query_attention: bool = True
     gated_linear_unit: bool = True
     activation_func: object = field(default=None)
+    # GDN (Gated DeltaNet) settings
+    experimental_attention_variant: str | None = None
+    linear_attention_freq: int | list | None = None
+    linear_conv_kernel_dim: int = 4
+    linear_key_head_dim: int = 128
+    linear_value_head_dim: int = 128
+    linear_num_key_heads: int = 16
+    linear_num_value_heads: int = 48
 
     def __post_init__(self):
         import torch.nn.functional as F
@@ -380,6 +388,204 @@ def test_swiglu_scaling_factor(self):
         assert flops_swiglu == expected_swiglu, f"SwiGLU: expected {expected_swiglu:.2e} but got {flops_swiglu:.2e}"
 
 
+class TestGDNLayerFlops:
+    """Tests for Gated DeltaNet (GDN) FLOPs calculation in transformer_flops path."""
+
+    def _qwen35_27b_config(self, **overrides):
+        """Return a MockModelConfig resembling Qwen3.5-27B (dense, 64 layers, freq=4)."""
+        defaults = dict(
+            num_layers=64,
+            hidden_size=5120,
+            seq_length=4096,
+            ffn_hidden_size=17408,
+            num_attention_heads=24,
+            num_query_groups=4,
+            kv_channels=256,
+            vocab_size=248320,
+            make_vocab_size_divisible_by=128,
+            tensor_model_parallel_size=1,
+            gated_linear_unit=True,
+            experimental_attention_variant="gated_delta_net",
+            linear_attention_freq=4,
+            linear_conv_kernel_dim=4,
+            linear_key_head_dim=128,
+            linear_value_head_dim=128,
+            linear_num_key_heads=16,
+            linear_num_value_heads=48,
+        )
+        defaults.update(overrides)
+        return MockModelConfig(**defaults)
+
+    def test_gdn_flops_differ_from_pure_attention(self):
+        """GDN-enabled config should produce different FLOPs than pure-attention baseline."""
+        batch_size = 1
+        gdn_cfg = MockConfigContainer(model=self._qwen35_27b_config())
+        baseline_cfg = MockConfigContainer(model=self._qwen35_27b_config(experimental_attention_variant=None))
+        gdn_flops = num_floating_point_operations(gdn_cfg, batch_size=batch_size)
+        baseline_flops = num_floating_point_operations(baseline_cfg, batch_size=batch_size)
+        assert gdn_flops != baseline_flops, "GDN FLOPs should differ from pure-attention FLOPs"
+        assert gdn_flops > 0
+
+    def test_gdn_only_layers(self):
+        """With linear_attention_freq=1 (no standard attn), self_attn_term should be pure GDN."""
+        batch_size = 1
+        num_layers = 4
+        hidden_size = 1024
+        seq_length = 512
+        vocab_size = 32000
+        qk_head_dim = 64
+        v_head_dim = 64
+        num_qk_heads = 8
+        num_v_heads = 16
+        conv_kernel_dim = 4
+
+        model_cfg = MockModelConfig(
+            num_layers=num_layers,
+            hidden_size=hidden_size,
+            seq_length=seq_length,
+            ffn_hidden_size=4096,
+            num_attention_heads=8,
+            num_query_groups=8,
+            kv_channels=128,
+            vocab_size=vocab_size,
+            make_vocab_size_divisible_by=128,
+            tensor_model_parallel_size=1,
+            gated_linear_unit=False,
+            experimental_attention_variant="gated_delta_net",
+            linear_attention_freq=1,
+            linear_conv_kernel_dim=conv_kernel_dim,
+            linear_key_head_dim=qk_head_dim,
+            linear_value_head_dim=v_head_dim,
+            linear_num_key_heads=num_qk_heads,
+            linear_num_value_heads=num_v_heads,
+        )
+        cfg = MockConfigContainer(model=model_cfg)
+        actual_flops = num_floating_point_operations(cfg, batch_size=batch_size)
+
+        # freq=1: pattern = [0 if (i+1)%1==0 else 1 for i in range(4)] = [0,0,0,0]
+        # All layers are standard attention, 0 GDN layers.
+        # This is because freq=1 means every layer is standard attention.
+        # So actual_flops should equal baseline (no GDN).
+        baseline_cfg = MockConfigContainer(
+            model=MockModelConfig(
+                num_layers=num_layers,
+                hidden_size=hidden_size,
+                seq_length=seq_length,
+                ffn_hidden_size=4096,
+                num_attention_heads=8,
+                num_query_groups=8,
+                kv_channels=128,
+                vocab_size=vocab_size,
+                make_vocab_size_divisible_by=128,
+                tensor_model_parallel_size=1,
+                gated_linear_unit=False,
+            )
+        )
+        baseline_flops = num_floating_point_operations(baseline_cfg, batch_size=batch_size)
+        assert actual_flops == baseline_flops, (
+            "freq=1 means every layer is standard attention, so FLOPs should match baseline"
+        )
+
+    def test_gdn_layer_freq_list(self):
+        """Test GDN with linear_attention_freq as a list pattern."""
+        batch_size = 1
+        # [1, 1, 0, 1, 1, 0, 1, 1] means 6 GDN, 2 standard
+        freq_list = [1, 1, 0, 1, 1, 0, 1, 1]
+        model_cfg = self._qwen35_27b_config(
+            num_layers=8,
+            linear_attention_freq=freq_list,
+        )
+        cfg = MockConfigContainer(model=model_cfg)
+        flops = num_floating_point_operations(cfg, batch_size=batch_size)
+        assert flops > 0
+
+    def test_gdn_exact_self_attn_term(self):
+        """Verify the GDN self_attn_term matches the expected formula from Megatron-LM."""
+        batch_size = 1
+        num_layers = 4
+        hidden_size = 1024
+        seq_length = 256
+        vocab_size = 32000
+        qk_head_dim = 64
+        v_head_dim = 64
+        num_qk_heads = 8
+        num_v_heads = 16
+        conv_kernel_dim = 4
+        ffn_hidden_size = 4096
+
+        qk_dim = qk_head_dim * num_qk_heads
+        v_dim = v_head_dim * num_v_heads
+
+        # freq=2: layers 0,2 are GDN (pattern[i]=1), layers 1,3 are standard (pattern[i]=0)
+        model_cfg = MockModelConfig(
+            num_layers=num_layers,
+            hidden_size=hidden_size,
+            seq_length=seq_length,
+            ffn_hidden_size=ffn_hidden_size,
+            num_attention_heads=8,
+            num_query_groups=8,
+            kv_channels=128,
+            vocab_size=vocab_size,
+            make_vocab_size_divisible_by=128,
+            tensor_model_parallel_size=1,
+            gated_linear_unit=False,
+            experimental_attention_variant="gated_delta_net",
+            linear_attention_freq=2,
+            linear_conv_kernel_dim=conv_kernel_dim,
+            linear_key_head_dim=qk_head_dim,
+            linear_value_head_dim=v_head_dim,
+            linear_num_key_heads=num_qk_heads,
+            linear_num_value_heads=num_v_heads,
+        )
+        cfg = MockConfigContainer(model=model_cfg)
+        gdn_flops = num_floating_point_operations(cfg, batch_size=batch_size)
+
+        # Compute expected manually
+        expansion_factor = 3 * 2 * 2  # 12
+        # Standard attention per-layer (MHA, num_query_groups==num_attention_heads so ratio=1)
+        kv_channels = 128
+        query_projection_size = kv_channels * 8
+        query_projection_to_hidden_size_ratio = query_projection_size / hidden_size
+        standard_attn_per_layer = (
+            expansion_factor
+            * hidden_size
+            * hidden_size
+            * ((1 + 8 / 8 + seq_length / hidden_size / 2) * query_projection_to_hidden_size_ratio)
+        )
+        # GDN per-layer
+        gdn_per_layer = (
+            3
+            * 2
+            * (
+                hidden_size * (2 * qk_dim + 2 * v_dim + 2 * num_v_heads)
+                + conv_kernel_dim * (2 * qk_dim + v_dim)
+                + num_v_heads * (v_head_dim**2) * 4
+                + hidden_size * v_dim
+            )
+        )
+        # freq=2: pattern = [1, 0, 1, 0] -> 2 GDN, 2 standard
+        expected_self_attn = gdn_per_layer * 2 + standard_attn_per_layer * 2
+        # MLP: gated_linear_unit=False -> gated_linear_multiplier=1
+        expected_mlp = expansion_factor * num_layers * hidden_size * ffn_hidden_size * 1
+        # Logit
+        padded_vocab = vocab_size  # 32000 is already divisible by 128
+        expected_logit = 3 * 2 * hidden_size * padded_vocab * 1
+        expected_total = batch_size * seq_length * (expected_mlp + expected_self_attn + expected_logit)
+
+        assert gdn_flops == expected_total, f"Expected {expected_total:.6e} but got {gdn_flops:.6e}"
+
+    def test_gdn_more_gdn_layers_changes_flops(self):
+        """Increasing GDN layer ratio (higher freq) should change FLOPs."""
+        batch_size = 1
+        # freq=4: 3/4 GDN, 1/4 standard
+        cfg_freq4 = MockConfigContainer(model=self._qwen35_27b_config(num_layers=8, linear_attention_freq=4))
+        # freq=8: 7/8 GDN, 1/8 standard
+        cfg_freq8 = MockConfigContainer(model=self._qwen35_27b_config(num_layers=8, linear_attention_freq=8))
+        flops_freq4 = num_floating_point_operations(cfg_freq4, batch_size=batch_size)
+        flops_freq8 = num_floating_point_operations(cfg_freq8, batch_size=batch_size)
+        assert flops_freq4 != flops_freq8, "Different GDN ratios should produce different FLOPs"
+
+
 class TestHybridMtpPatternParsing:
     """Tests for hybrid/MTP pattern parsing in FLOPs accounting."""