[split 4/4] Enable DSA CP and THD hooks - #5246
Conversation
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 4245d4d58b
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Provide absorbed MLA submodule names
With experimental_attention_variant == "dsa", this now instantiates AbsorbedMLASelfAttentionSubmodules, but the spec still passes the old MLA field linear_kv_up_proj. That dataclass defines linear_k_up_proj and linear_v_up_proj instead, so constructing the DSA module spec raises an unexpected-keyword TypeError before any model can be built in this path. Please split the absorbed K/V up projections into the field names expected by AbsorbedMLASelfAttentionSubmodules.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
888e0ec to
2553904
Compare
There was a problem hiding this comment.
💡 Codex Review
Use the absorbed MLA submodule names
With this switch to AbsorbedMLASelfAttentionSubmodules, constructing the DSA module spec now passes linear_kv_up_proj, but that dataclass only accepts linear_k_up_proj and linear_v_up_proj. As a result, any caller of get_dsa_module_spec_for_backend() with experimental_attention_variant="dsa" fails immediately with an unexpected-keyword TypeError before the model can be built.
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
2553904 to
1042876
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 181710c61c
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 7062b3878a
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _BACKEND_MODULE_NAME_BY_BACKEND = { | ||
| "tilelang": "megatron.core.transformer.experimental_attention_variant.dsa_tilelang_kernels", | ||
| "cudnn": "megatron.core.transformer.experimental_attention_variant.dsa_cudnn_kernels", | ||
| } |
There was a problem hiding this comment.
When users set dsa_kernel_backend to tilelang or cudnn, use_fused_dsa_kernels() returns true and the first fused hook calls _load_backend(), but these module names do not exist in this tree (I checked repo-wide with rg --files for dsa_tilelang_kernels, dsa_cudnn_kernels, and tilelang_dsa). In that configuration the forward path raises RuntimeError: Failed to import DSA kernel backend ... instead of using the advertised optional backend or falling back, so either include/register these backend modules here or keep those choices disabled until they are available.
Useful? React with 👍 / 👎.
|
/claude review |
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
/ok to test 8f75c5c |
There was a problem hiding this comment.
Can you explain why the changes in this file are necessary? In particular, why do we need to manually construct this checkpoint_inputs list and specifically overwrite the attn_mask_type value?
There was a problem hiding this comment.
Added code comments for the explanation. Some more details here:
The change is needed for activation checkpointing (recompute) when the core attention needs extra tensor inputs, specifically the DSA path in this case. In the DSA/MLA flow, core_attention_extra_kwargs can contain tensors like x=hidden_states and qr=q_compressed. Those tensors are used inside core_attention and must be explicit inputs to tensor_parallel.checkpoint; otherwise they are captured from the Python closure, which can either defeat checkpointing by retaining the original forward tensors or give incorrect/backward-fragile behavior because checkpoint recompute only detaches and returns gradients for tensors passed as checkpoint inputs.
attn_mask_type is handled specially because it is metadata semantically, but Megatron’s checkpoint wrapper saves its *args via ctx.save_for_backward(*args), so the checkpoint argument list must contain tensors, not an AttnMaskType enum. We preserve the existing behavior from main by encoding the enum value as a small tensor before calling checkpoint, then converting it back inside custom_forward with AttnMaskType(attn_mask_type.item()).
The overwrite of checkpoint_inputs[5] is just because the list is initially constructed with a placeholder before defaulting attn_mask_type and converting it to the tensor form. It is not changing the mask type; it replaces None/enum metadata with the tensor representation required by checkpointing.
| if seq_lens.numel() == 0: | ||
| return torch.empty(0, dtype=torch.int64, device=device) | ||
|
|
||
| if cu_seqlens_i64.device.type == "cpu": |
There was a problem hiding this comment.
What does this block do? IIUC the sequence lengths are only known at runtime and vary per batch, so divisibility can only really be guaranteed by the data pipeline's padding and this block is a no-op during training, is that right?
There was a problem hiding this comment.
Yes, this block is a host-side validation guard for CPU/test/debug paths. Packed zigzag CP assumes each packed sequence is padded to a multiple of 2 * cp_size, so the rank-local front/back chunks are integral. In training we avoid checking CUDA runtime lengths here to prevent synchronization; the data pipeline is responsible for that padding guarantee. My preference is to keep it with the comment because it documents and validates the packed zigzag CP invariant without affecting CUDA training, but we can also get this removed if you prefer.
|
|
||
|
|
||
| def build_causal_mask_from_positions( | ||
| query_pos: torch.Tensor, key_pos: torch.Tensor |
There was a problem hiding this comment.
Q: it looks like key_pos is always an arange, is there a use case where key_pos is shuffled?
There was a problem hiding this comment.
In the current main path, after gathered KV is restored to global order, key_pos is usually arange. I kept it explicit because this helper is a public function that meant to build masks from logical token positions rather than assuming KV storage order. That also covers cases where we build/apply the mask before KV reorder, or if a caller passes subset/reordered KV.
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
Signed-off-by: Hollow Man <hollowman@opensuse.org>
|
/ok to test f58de5b |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/28195685815 |
6 participants
Split out from #3674 to reduce review scope.
Original changes by @HollowMan6 in #3674.
Scope
multi_latent_attention.py.Dependencies
CODEOWNERS note
megatron/training/arguments.pyhas an exact CODEOWNERS entry with no owners after the generalmegatron/training/rule, so this PR includes that one-line skip-field update here and calls out the ownership ambiguity explicitly.Validation
git diff --check upstream/main..split-3674-dsa-cp-thd-core