From c6daa546d8b7b06b09a8e40534d0c30eba28cf29 Mon Sep 17 00:00:00 2001 From: Vasudevan Rengasamy Date: Fri, 28 Aug 2026 00:04:40 -0700 Subject: [PATCH] Add gdn_kernel_backend knob Signed-off-by: Vasudevan Rengasamy --- nemo_rl/models/megatron/setup.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/nemo_rl/models/megatron/setup.py b/nemo_rl/models/megatron/setup.py index 399633cea52..fa983592d6b 100644 --- a/nemo_rl/models/megatron/setup.py +++ b/nemo_rl/models/megatron/setup.py @@ -1121,6 +1121,8 @@ def _apply_performance_config(model_cfg: Any, config: PolicyConfig) -> None: # Flash attention does not support the resulting padded multi-row layout, # so the canonical expanded-sequence contract must use backend dispatch. attention_backend = config["megatron_cfg"].get("attention_backend") + if "gdn_kernel_backend" in config["megatron_cfg"]: + model_cfg.gdn_kernel_backend = config["megatron_cfg"]["gdn_kernel_backend"] if ( getattr(model_cfg, "nemotron_omni_contract", None) == _NEMOTRON_OMNI_EXPANDED_SEQUENCE_CONTRACT