fix: attach CP attention-mask hooks for dense (non-TE) context parallelism#1470

Open

hemildesai wants to merge 11 commits intomainfrom

hemil/cp-dense-fixes

Commits on Mar 9, 2026

fix: attach CP attention-mask hooks for dense (non-TE) context parallelism

hemildesai
and
claude
committed
fix: use late-bound SDPA lookup for context parallelism compatibility

hemildesai
and
claude
committed
feat: add resolve_sdpa_method helper with YAML-configurable SDPA backends

hemildesai
and
claude
committed
style: apply ruff formatting to changed files

hemildesai
and
claude
committed
fix: skip non-TE attention in MoE apply_cp instead of asserting

hemildesai
and
claude
committed
refactor: move resolve_sdpa_method to model_init.py

hemildesai
and
claude
committed
refactor: move CP hook attachment into apply_model_infrastructure

hemildesai
and
claude
committed
fix: check model parts in _uses_te_attention for AutoPipeline support

hemildesai
and
claude
committed
refactor: call resolve_sdpa_method inside model init, not train_ft

hemildesai
and
claude
committed
fix: warn when sdpa_method is ignored for non-NeMoAutoModel targets

hemildesai
and
claude
committed
fix: update softmax_scale tests for closure-based SDPA attn_func

hemildesai
and
claude
committed