Naver Cloud Achieves 1.33x Speedup on Single Layer via Mamba-2 Kernel Optimization for B200 GPUs

Naver Cloud optimized the Mamba-2 kernel to maximize the hardware potential of NVIDIA Blackwell (B200) GPUs. The existing open-source kernel, mamba-ssm, failed to fully utilize the new hardware features of the B200 due to heuristic settings based on the A100, prompting the adoption of Helion DSL to address this issue.

Three-Stage Optimization Process Based on Helion

The optimization process proceeded in three stages: bottleneck analysis, numerical precision assurance, and logic and heuristic improvement. First, the top 10 kernels accounting for 82% of the total computation time were profiled to identify bottlenecks. Next, the mathematical computation mechanism was reproduced 1:1 with the existing Triton implementation to minimize numerical errors. Finally, the computation structure was reconfigured to reflect the architectural characteristics of the B200. In particular, the sequential loop of the _state_passing operation was converted to the Associative Scan command based on Helion to maximize parallelism.

Bottleneck Resolution and Performance Improvement

The _layer_norm_bwd kernel, one of the core bottlenecks, caused performance degradation because fixed parameters based on the A100 did not reflect the increased bandwidth and latency imbalance of the B200. By using Helion's Autotuner to precisely search and adjust the nrow_groups parameter from the existing 74 to 2601 to match the B200 structure, a 3.18x speedup was achieved in that kernel. This result was obtained by recovering the potential of the existing hardware through parallelism parameter adjustments following the hardware generation transition, rather than introducing a new algorithm.

Application to Actual Training Environments and Effects

To reduce autotuning overhead, an Automated Config Generator was built to pre-save and load optimal parameters for each training environment. This allows optimized performance to be immediately utilized in the actual Pre-training stage without 1-2 hours of waiting time. In the End-to-End training test of the Nemotron-3 Nano 30B model, the processing speed per total training step was reduced by 1.12x, and the Loss convergence curve showed the same pattern as the Baseline, proving stability. Naver Cloud plans to improve the overall computation efficiency of the cluster by linking with distributed training paradigms such as Context Parallelism in the future.