Dev - #406
Conversation
Summary of ChangesHello @xiuhu17, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request significantly enhances the DeepSeekV32 model's attention mechanism by integrating Context Parallelism (CP) support for sparse Multi-Latent Attention (MLA). The changes involve a refactoring of the attention backward pass to handle combined Key-Value (KV) tensors, alongside a more flexible approach to generating and applying attention masks in a distributed environment. A new wrapper leverages optimized sparse MLA kernels, and a dedicated test ensures the functional correctness of these context-parallel optimizations. This work enables the DeepSeekV32 model to scale more efficiently by distributing attention computations across multiple devices. Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces significant changes to enable context parallelism for Dynamic Sparse Attention (DSA), including new and updated patch files for both the Megatron core and custom TileLang kernels. The changes are extensive and add complex logic for distributed attention computation.
While the effort to add context parallelism is appreciated, I've found several critical issues that need to be addressed:
- There's a critical bug in the backward pass of the attention mechanism in both
cp_kernel/megatron.patchandcp_kernel/tilelang.patchwhere gradients for key and value vectors are handled incorrectly. This will lead to incorrect model training. - A class was renamed in one file but not in the file that uses it, which will cause a runtime import error.
- The code contains hardcoded 'magic numbers' for tensor dimensions, which makes it fragile and hard to maintain.
- There are inconsistencies between the Python wrapper code and the TileLang kernels regarding causal masking.
These issues are detailed in the specific comments. Addressing them is crucial for the correctness and maintainability of the code. In the future, please provide a more descriptive title and a summary in the pull request body to help reviewers understand the context of the changes.
|
|
||
|
|
||
| -class AttentionFuncionWithContextParallel(torch.autograd.Function): | ||
| +class Ref(torch.autograd.Function): |
There was a problem hiding this comment.
This class was renamed to Ref, but it's still imported and used as AttentionFuncionWithContextParallel in megatron/core/transformer/experimental_attention_variant/dsa.py (also part of this patch). This will cause an ImportError at runtime. Please use a consistent name across files.
Additionally, Ref is not a descriptive name. I'd recommend reverting to AttentionFuncionWithContextParallel or choosing a more meaningful name.
class AttentionFuncionWithContextParallel(torch.autograd.Function):
|
|
||
| # Initialize or resume constants and communication group | ||
| - q, k, v, attention_mask, *rest = ctx.saved_tensors | ||
| + q, kv, _, attention_mask, *rest = ctx.saved_tensors |
There was a problem hiding this comment.
There's a critical bug in how k and v tensors are handled in the backward pass. The forward pass saves k and v as separate tensors via ctx.save_for_backward(q, k, v, ...). However, here in the backward pass, you unpack them as q, kv, _, ..., treating the saved k as a combined kv tensor and completely discarding the saved v. The subsequent logic, especially in eager_attn_bwd, assumes v can be derived from k (as kv), which is not true for the general case, for instance, when called from unfused_dsa_fn_with_cp in dsa.py with separate key and value tensors. This will lead to incorrect gradient computation for v and k. You should handle k and v as separate tensors throughout the backward pass, consistent with the forward signature.
| + '''Backward pass for the native attention function with context parallelism''' | ||
| + | ||
| + # Initialize or resume constants and communication group | ||
| + q, kv, _, attention_mask, *rest = ctx.saved_tensors |
There was a problem hiding this comment.
There's a critical bug in how k and v tensors are handled in the backward pass. The forward pass saves k and v as separate tensors, but here you unpack them as q, kv, _, ..., effectively treating the saved k as a combined kv tensor and discarding the saved v. The subsequent logic assumes v can be derived from k, which is not guaranteed by the forward signature. This will lead to incorrect gradient computation. The test case in test.py happens to satisfy this assumption, but the function itself is not safe for general use where k and v might be independent.
| + v = kv[:,:,:,:512] | ||
| + q_tail = q[:,:,:,512:] |
There was a problem hiding this comment.
Using hardcoded values like 512 makes the code fragile. These dimensions should be derived from tensor shapes dynamically. For example, you could get the value dimension from v.shape[-1] (though v is not available here due to another issue) or pass it as an argument. The same applies to 576 used later, which seems to be 512 + 64. Please replace these magic numbers with dynamically determined values or constants defined elsewhere.
| + v = kv[:,:,:,:512] | ||
| + q_tail = q[:,:,:,512:] |
There was a problem hiding this comment.
Using hardcoded values like 512 makes the code fragile and hard to maintain. These dimensions should be derived from tensor shapes dynamically or passed as arguments. The same applies to 576 used later. Please replace these magic numbers with dynamically determined values or constants.
| # Compute attention scores | ||
| for h_i, bi_i in T.Parallel(block_H, BS): | ||
| - acc_p[h_i, bi_i] = T.if_then_else(mask[bi_i], 0, -T.infinity(acc_p.dtype)) | ||
| + acc_p[h_i, bi_i] = 0 |
There was a problem hiding this comment.
The causal masking logic has been removed from this kernel (and from the forward kernel as well) by initializing acc_p to 0. However, the Python wrapper sparse_mla_bwd in ring_wrapper.py is called with is_causal=True, and the wrapper function in sparse_mla_bwd.py still has logic depending on this flag. This is inconsistent and misleading. If causality is no longer supported by the kernel, the is_causal parameter and related logic should be removed from the Python wrappers to avoid confusion and potential bugs.
2 participants
No description provided.