[KDA] Add Paddle training compatibility - #2
Merged
Conversation
SigureMo
reviewed
Jul 30, 2026
SigureMo
left a comment
SigureMo
left a comment
There was a problem hiding this comment.
预期，如下代码是可以跑通的
import paddle
paddle.enable_compat(scope={'fla', 'triton'})
import fla
paddle.disable_compat()
from fla.modules import FusedRMSNormGated, ShortConvolution # noqa: E402
from fla.ops.kda import chunk_kda # noqa: E402
from fla.ops.utils.index import prepare_cu_seqlens_from_mask, prepare_lens_from_mask # noqa: E402
from fla.utils import tensor_cache # noqa: E402
# use FusedRMSNormGated, ShortConvolution, chunk_kda, prepare_cu_seqlens_from_mask, prepare_lens_from_mask, tensor_cache
Comment on lines
+152
to
+167
| """Return the default implementation without optional backend routing. | ||
|
|
||
| This branch supports KDA training on NVIDIA GPUs through the default Triton implementation. | ||
| The optional KDA backends are: | ||
|
|
||
| - TileLang is deferred for now. | ||
| TODO: consider supporting it if performance is better than the Triton implementation. | ||
| - FlashKDA, which is inference-only and has no training backward. | ||
| - Triton-Ascend, which only targets Huawei NPUs. | ||
|
|
||
| The module registry similarly contains only Triton-Ascend overrides. | ||
| None of these backends are currently required for the supported CUDA training paths. | ||
| Those paths include dense and variable-length forward/backward and context parallelism. | ||
| Loading them would only expand the unsupported dependency and import surface. | ||
|
|
||
| Iterates through all registered backends and selects the first one | ||
| that passes the verifier for the given function call. | ||
| The upstream dispatch wrapper depends on ``torch.compiler.disable``, which Paddle compat does not provide. |
Comment on lines
-167
to
-186
| if _DISPATCH_DISABLED: | ||
| return func | ||
| func_name = func.__name__ | ||
|
|
||
| @wraps(func) | ||
| def wrapper(*args, **kwargs) -> Any: | ||
| # Lazy initialization of backends | ||
| BackendRegistry.ensure_initialized(operation) | ||
|
|
||
| registry = BackendRegistry._registries.get(operation) | ||
| if registry is None: | ||
| return func(*args, **kwargs) | ||
|
|
||
| # Iterate through all registered backends sorted by priority | ||
| # to find one that can handle this call | ||
| backends_list = registry._get_sorted_backends() | ||
|
|
||
| for be in backends_list: | ||
| # Avoid be.can_use(): its @cache wrapper breaks torch.compile tracing. | ||
| if not (be.is_available() and be.is_enabled()): |
| return output, None | ||
|
|
||
| if backend == 'triton': | ||
| from fla.modules.conv.triton import CausalConv1dFunction |
SigureMo
approved these changes
Jul 30, 2026
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Add a compact Paddle compatibility path for KDA training while retaining the existing Triton kernels and shared FLA implementations. The change supports dense and variable-length
chunk_kdaforward/backward, synchronous context parallelism, and theShortConvolutionandFusedRMSNormGateddependencies used by the KimiDeltaAttention training data flow.Paddle compatibility is enabled explicitly by the consuming process before importing FLA. The library itself does not call
paddle.enable_compat, andfla.ops.kdaexposes onlychunk_kda.Test plan
tests/paddle/test_kda.pywith dense save/recompute, gate, beta, state, gradient, GVA, and variable-length coverage against a Paddle eager recurrence.tests/paddle/test_cp_kda.pywith two-GPU context-parallel forward/backward checks.tests/paddle/test_kimi_delta_attention.pyfor ShortConvolution, gated RMSNorm, import boundaries, and dense/masked KimiDeltaAttention training data flow.CUDA_VISIBLE_DEVICES=0,1 ../.venv/bin/python -m pytest --confcutdir=tests/paddle tests/paddle -q: 12 passed.pre-commit run --files <changed files>: all hooks passed.Benchmark / NCU
Not run. This is a framework compatibility change and does not alter the KDA Triton kernel algorithm, tiling, or scheduling.
Breaking changes
This target branch is Paddle-only and is not intended to preserve the complete upstream PyTorch/FLA API:
flaandtritonscopes before importing FLA;fla.ops.kdaexportschunk_kdaonly;