[megatron] 3/n towards Kimi K2.6: skip MLA THD value pad on sm100+ to keep fused attention trainable

skyrl/backends/skyrl_train/patches/megatron/patch_mla_thd_v_pad.py

@@ -0,0 +1,56 @@
+"""Skip Megatron's MLA THD value pad on Blackwell (sm100+).
+
+For packed (``qkv_format="thd"``) execution, Megatron-core pads the MLA value
+tensor from ``v_head_dim`` (e.g. 128) up to the QK head dim (e.g. 192) in
+``_prepare_mla_core_attention_value`` and trims the attention output back
+afterwards.
+
+On Blackwell that pad is fatal for training: cuDNN fused attention has no
+backward support for ``head_dim > 128`` on sm100+, so with the padded
+``head_dim_v == head_dim_qk == 192`` TransformerEngine disables FusedAttention
+for training-mode forwards. FlashAttention 2 does not support MLA at all,
+FlashAttention 3 is sm90-only, and UnfusedDotProductAttention does not support
+context parallelism - so MLA + CP training raises
+``ValueError: No dot product attention backend is available``. Inference-mode
+forwards (logprob computation) are unaffected, which makes the failure appear
+only at the first ``forward_backward``.
+
+cuDNN fused attention natively supports MLA's unequal QK/V head dims
+(192/128), including THD + context parallelism with the ``p2p`` exchange, for
+both forward and backward. Skipping the pad simply selects that native path
+(and saves the pad/trim memory traffic). Behavior on pre-Blackwell devices is
+left unchanged.
+"""
+
+from loguru import logger
+
+_APPLIED = False
+
+
+def patch_mla_thd_v_pad() -> None:
+    """Patch ``_prepare_mla_core_attention_value`` to skip the V pad on sm100+."""
+    global _APPLIED
+    if _APPLIED:
+        return
+
+    import torch
+    from megatron.core.transformer import multi_latent_attention as mla
+
+    orig_prepare = mla._prepare_mla_core_attention_value
+
+    def patched_prepare(parallel_attention, query, value, packed_seq_params):
+        if (
+            value is not None
+            and packed_seq_params is not None
+            and getattr(packed_seq_params, "qkv_format", None) == "thd"
+            and query.shape[-1] != value.shape[-1]
+            and torch.cuda.is_available()
+            and torch.cuda.get_device_capability() >= (10, 0)
+        ):
+            orig_v_dim = value.shape[-1]
+            return value, False, orig_v_dim, orig_v_dim
+        return orig_prepare(parallel_attention, query, value, packed_seq_params)
+
+    mla._prepare_mla_core_attention_value = patched_prepare
+    _APPLIED = True
+    logger.info("Applied Megatron MLA THD V-pad skip for sm100+ (native unequal-head-dim attention)")

skyrl/backends/skyrl_train/workers/megatron/megatron_worker.py

@@ -19,6 +19,7 @@
 from omegaconf import OmegaConf
 from transformers import AutoConfig
 
+import skyrl.backends.skyrl_train.workers.megatron.model_bridges  # noqa: F401  # register extra bridges
 from skyrl.backends.skyrl_train.distributed.dispatch import MeshRank, WorkerOutput
 from skyrl.backends.skyrl_train.distributed.megatron.megatron_strategy import (
     MegatronStrategy,
@@ -44,6 +45,9 @@
 from skyrl.backends.skyrl_train.patches.megatron.patch_dsa_index_share import (
     patch_dsa_index_share,
 )
+from skyrl.backends.skyrl_train.patches.megatron.patch_mla_thd_v_pad import (
+    patch_mla_thd_v_pad,
+)
 from skyrl.backends.skyrl_train.patches.te.patch_fa2_head_dim import (
     patch_fa2_head_dim_allowlist,
 )
@@ -66,6 +70,9 @@
 from skyrl.backends.skyrl_train.workers.megatron.megatron_model_wrapper import (
     MegatronModelWrapper,
 )
+from skyrl.backends.skyrl_train.workers.megatron.model_bridges import (
+    maybe_force_qwen35_text_bridge,
+)
 from skyrl.backends.skyrl_train.workers.worker import (
     CriticWorkerBase,
     PolicyWorkerBase,
@@ -84,17 +91,14 @@
 from skyrl.train.utils.utils import str_to_torch_dtype, update_model_config
 from skyrl.utils.tok import get_tokenizer
 
+patch_mla_thd_v_pad()
+
 if TYPE_CHECKING:
     from skyrl.backends.skyrl_train.inference_servers.base import (
         InferenceEngineInterface,
     )
     from skyrl.train.config.config import InferenceEngineConfig
 
-import skyrl.backends.skyrl_train.workers.megatron.model_bridges  # noqa: F401  # register extra bridges
-from skyrl.backends.skyrl_train.workers.megatron.model_bridges import (
-    maybe_force_qwen35_text_bridge,
-)
-
 
 class MegatronWeightExtractor(WeightExtractor):
     """Extracts weights from Megatron model-parallel models.

skyrl/train/utils/utils.py

@@ -741,6 +741,13 @@ def prepare_runtime_environment(cfg: SkyRLTrainConfig) -> dict[str, str]:
             # https://github.com/NVIDIA/TransformerEngine/blob/release_v2.5/transformer_engine/pytorch/attention/dot_product_attention/utils.py#L916
             env_vars["NVTE_FUSED_ATTN"] = "0"
 
+        # Forward TransformerEngine attention-backend debug logging to workers when
+        # set on the driver. Workers are re-exec'd through the runtime env (e.g. the
+        # uv hook), so a plain raylet/driver export does not reach them.
+        for nvte_var in ("NVTE_DEBUG", "NVTE_DEBUG_LEVEL"):
+            if os.environ.get(nvte_var):
+                env_vars[nvte_var] = os.environ[nvte_var]
+
     if cfg.generator.inference_engine.backend == "vllm":
         env_vars["VLLM_ALLOW_RUNTIME_LORA_UPDATING"] = "true"