Add GDN dynamic inference - #6499
Conversation
952025b to
3c93c80
Compare
|
/ok to test 3c93c80 |
| # Mamba and GDN use the same slot-indexed recurrent-state cache contract. Build | ||
| # one map in global layer order; independently generated per-symbol maps both | ||
| # start at zero and would alias if they were simply unioned. | ||
| attention_layer_map, dsa_layer_map = operator.itemgetter( | ||
| Symbols.ATTENTION, Symbols.DS_ATTENTION | ||
| )(get_layer_maps_from_layer_type_list(mamba_inference_state_config.layer_type_list)) | ||
| recurrent_layer_map = {} | ||
| for global_layer_idx, layer_type in enumerate( | ||
| mamba_inference_state_config.layer_type_list | ||
| ): | ||
| if layer_type in (Symbols.MAMBA, Symbols.GDN): | ||
| recurrent_layer_map[global_layer_idx] = len(recurrent_layer_map) |
There was a problem hiding this comment.
Not really related to this PR, but I'm curious why GDN has a specific layer type but GDP did not - should we be unifying all of these linear attention variants under a single layer type?
There was a problem hiding this comment.
good question! I don’t know why GDP doesn’t have its own symbol. Symbol.GDN was already there and being used, so it made sense to just follow the pattern for this work. GDP piggybacks on Symbol.Mamba, I think.
| @@ -0,0 +1,180 @@ | |||
| # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. | |||
|
|
|||
| """Focused kernel and CUDA-graph tests for Gated DeltaNet dynamic inference.""" | |||
There was a problem hiding this comment.
Is cuda graph functionality intended to work out of the box? For Mamba2 at least we needed to modify the kernel to support masking computation when we pass -1 values in the batch_indices tensor, do the GDN kernels already support this?
There was a problem hiding this comment.
writes use tensor_masked_update, which explicitly ignores -1.
for reads, there is a call to batch_indices.clamp(min=0) before gathering. this is present for both ssm_decode() and ssm_prefill() inside gdn.py.
There was a problem hiding this comment.
I see, so this is basically relying on the fact that multiple readers of the value at index 0 is safe as long as we don't have multiple writers. I think that's ok for now, but we should maybe document this explicitly.
There was a problem hiding this comment.
as long as we don't have multiple writers
@santhnm2, to clarify: writes go through tensor_masked_update, which in turn calls _tensor_masked_update_kernel_2d (or 3d/4d). these kernels ignore an index of -1 with an explicit return.
if target_idx == -1:
return
(it is possible I misunderstood your comment)
There was a problem hiding this comment.
Sorry let me clarify - my understanding is the clamp on batch_indices is basically forcing all padding entries to read from slot 0 and then do some dummy computation based on that value, but the result of the dummy computation is not written back to slot 0 because tensor_masked_update prevents writes on -1 slots. So what I meant by not having multiple writers is slot 0 will have a max of 1 writer (i.e., a real request mapped to slot 0) so this is safe.
The semantics for Mamba2 (and eventually GDP) are that within the kernel itself we skip computation for the -1 indices and directly write out 0 values rather than proceeding with the dummy computation. But the end result is effectively the same other than the output activation tensor having nonzero values in the padding slots. These padding slot outputs are ignored anyway so there should be no downstream effect.
There was a problem hiding this comment.
I get it now. the way Mamba2 handles padding sounds cleaner. probably worth changing this in another PR.
There was a problem hiding this comment.
Was this change intentional?
There was a problem hiding this comment.
intentional. FLA was previously a dev dependency. that seemed wrong as it is a requirement for GDN (just like causal-conv1d, which was already under the ssm section). but let me know if you think FLA belongs in dev.
There was a problem hiding this comment.
It seems like we should bump this, but maybe @NVIDIA/mcore-oncall has thoughts here?
Integrate GDN with the shared recurrent-state cache and FLA kernels for packed prefill and decode. Add configuration validation plus state-continuity coverage. Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Anil Thomas <anthomas@nvidia.com>
3c93c80 to
b6ff5f4
Compare
|
/ok to test b6ff5f4 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32435076560 |
|
🔄 Merge queue validation started! You can track the progress here: https://github.com/NVIDIA/Megatron-LM/actions/runs/32440023874 |
Signed-off-by: Anil Thomas <anthomas@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Dmytro Pykhtar <dpykhtar@nvidia.com>
Signed-off-by: Anil Thomas <anthomas@nvidia.com> Co-authored-by: OpenAI Codex <codex@openai.com> Signed-off-by: Kezhi Kong <kezhik@kezhik-mlt.client.nvidia.com>
5 participants
What does this PR do?
This implementation adds GDN prefill and single-token decode to the dynamic inference engine. It reuses the generic recurrent-mixer control flow introduced by PR #5382 and uses the inference kernels from flash-linear-attention (FLA) 0.5.1.
Contribution process
Pre-checks
Notes
For a detailed description of the changes, please see Gated DeltaNet Dynamic Inference (internal document)