What does this PR do ?

PR for main #3042

Changes

1. Switch to FLA conv1d. Benefits:
   • FLA conv1d accepts [b, s, h] shaped input, which helps us eliminate 1 transpose op.
   • FLA conv1d supports cu_seqlens arg, which is necessary for one of our future works, sequence packing.
   • FLA conv1d supports initial states, which is necessary for one of our future works, chunked PP.
2. Combine l2norm to query and key to eliminate 1 × elementwise_kernel.

E2E correctness is checked as follows. (baseline is in pink, the new impl is in blue)

Perf gain

Model: Qwen3-Next-80B-A3B

E2E perf gain: ~1.04x

Timelines:

Baseline GDN forward (1.79 ms, redundant ops are in red boxes)

Optimized GDN forward (1.49 ms)

Baseline GDN backward (3.24 ms, the redundant op is in the red box)

Optimized GDN backward (3.01 ms)

Contribution process

flowchart LR
    A[Pre-checks] --> B[PR Tests]
    subgraph Code Review/Approval
        C1[Expert Review] --> C2[Final Review]
    end
    B --> C1
    C2 --> D[Merge]

Pre-checks

• I want this PR in a versioned release and have added the appropriate Milestone (e.g., Core 0.8)
• I have added relevant unit tests
• I have added relevant functional tests
• I have added proper typing to my code Typing guidelines
• I have added relevant documentation
• I have run the autoformatter.sh on my PR

Code review

The following process is enforced via the CODEOWNERS file for changes into megatron/core. For changes outside of megatron/core, it is up to the PR author whether or not to tag the Final Reviewer team.

Merging your PR

Any member of core-adlr and core-nemo will be able to merge your PR.