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

[casper-hansen]

What

Skips Megatron-core's MLA THD value pad (128 -> 192) on Blackwell (sm100+) so TransformerEngine selects cuDNN fused attention for training-mode forwards. Also forwards NVTE_DEBUG / NVTE_DEBUG_LEVEL from the driver to Ray workers so TE attention-backend selection can actually be debugged (raylet/driver exports do not survive the runtime-env re-exec).

Why

cuDNN has no backward support for head_dim > 128 on sm100+, so the padded dims disable FusedAttention for training. With FlashAttention unavailable for MLA and unfused attention unsupported under context parallelism, the first forward_backward raised "No dot product attention backend is available" (inference-mode logprob forwards were unaffected, which made this training-only). cuDNN natively handles MLA's unequal head dims (192/128) for fwd+bwd including THD + CP p2p, so skipping the pad selects the native path and saves the pad/trim traffic. Pre-Blackwell behavior is unchanged.

Verified on 2x8 B300: Kimi K2.7 GRPO LoRA smoke completes training steps with the fused backend selected for is_training=True; Kimi bridge GPU tests pass 3/3.

Part of the Kimi K2.x series (follow-up to #1862). Independent of the other PRs in the series.

Made with Cursor

[gemini-code-assist]

Code Review

This pull request introduces a patch to skip Megatron's MLA THD value padding on Blackwell (sm100+) devices, preventing training failures caused by unsupported head dimensions in cuDNN fused attention. It also forwards TransformerEngine debug environment variables to workers. The review feedback suggests making the monkey-patching more robust by checking for the existence of the private Megatron method before patching, and guarding the CUDA capability check with torch.cuda.is_available() to prevent crashes in CPU-only environments.

[gemini-code-assist]

Monkey-patching private methods of external libraries (like Megatron-core's _prepare_mla_core_attention_value) can be fragile across library updates. If the method is renamed or removed in a future version, importing this module will raise an AttributeError and crash the application. It is safer to check for the existence of the attribute before patching it.

    orig_prepare = getattr(mla, "_prepare_mla_core_attention_value", None)
    if orig_prepare is None:
        logger.warning("Megatron MLA THD V-pad patch skipped: _prepare_mla_core_attention_value not found.")
        return

[casper-hansen]

Deliberately not guarding this one: megatron-core is pinned to an exact revision (uv.lock / the deploy manifest), so the symbol can only disappear on an intentional pin bump — and in that case we want a loud AttributeError at import. Skipping the patch with a warning would instead resurface the failure this patch exists to fix ("No dot product attention backend is available" at the first forward_backward on sm100+), which is far harder to trace back to a missing patch.

[gemini-code-assist]

Calling torch.cuda.get_device_capability() directly can raise an exception (e.g., AssertionError or RuntimeError) in CPU-only environments, such as local testing or certain CI/CD pipelines where CUDA is not compiled or no GPU is available. Guarding this call with torch.cuda.is_available() prevents unexpected crashes.

            and torch.cuda.is_available()
            and torch.cuda.get_device_capability() >= (10, 0)

[casper-hansen]

Done in 2f72c7b — CPU-only environments now fall through to the original pad path (the skip only matters where cuDNN fused attention runs).