KDA attention - #3548
KDA attention#3548ksivaman wants to merge 3 commits into
Conversation
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
Signed-off-by: Kirthi Shankar Sivamani <ksivamani@nvidia.com>
for more information, see https://pre-commit.ci
 |
![greptile-apps[bot]](https://avatars.githubusercontent.com/in/867647?s=60&v=4){kind=small}
| _KDA_OP_AVAILABLE = _kda_op_available() | ||
| if os.getenv("NVTE_KDA_TEST_REQUIRED", "0") == "1" and not _KDA_OP_AVAILABLE: | ||
| raise RuntimeError( | ||
| "NVTE_KDA_TEST_REQUIRED=1, but the cuDNN frontend KDA op or its " | ||
| "cutedsl runtime is unavailable." | ||
| ) |
There was a problem hiding this comment.
If the L0 suite runs on a non-SM100/SM103/SM107 GPU that has the KDA op and CUTLASS installed, NVTE_KDA_TEST_REQUIRED=1 does not raise, but the requires_kda marker skips every KDA test. The command can therefore succeed without exercising the feature, giving false assurance that KDA was tested. Required mode should also fail when the current architecture cannot run KDA.
| if allow_neg_eigval and not use_beta_sigmoid_in_kernel: | ||
| raise ValueError( | ||
| "KimiDeltaAttention allow_neg_eigval requires " | ||
| "use_beta_sigmoid_in_kernel, which owns the sigmoid it scales." | ||
| ) | ||
| if gate_domain not in {"log", "linear"}: | ||
| raise ValueError( | ||
| f"KimiDeltaAttention gate_domain must be 'log' or 'linear', got {gate_domain!r}." | ||
| ) | ||
| if not safe_gate: | ||
| extra = [ | ||
| name | ||
| for name, value in ( | ||
| ("a_log", a_log), | ||
| ("dt_bias", dt_bias), | ||
| ("gate_lower_bound", gate_lower_bound), | ||
| ) | ||
| if value is not None | ||
| ] | ||
| if extra: | ||
| raise ValueError( | ||
| f"KimiDeltaAttention {', '.join(extra)} require safe_gate, which owns " | ||
| "the transform they parameterize." | ||
| ) | ||
| elif gate_domain == "linear": | ||
| raise ValueError( | ||
| "KimiDeltaAttention gate_domain='linear' cannot combine with safe_gate; " | ||
| "the safe-gate transform takes raw logits." | ||
| ) |
There was a problem hiding this comment.
The new public validation branches have no negative tests. This leaves contracts such as rejecting allow_neg_eigval without fused sigmoid, invalid gate domains, safe-gate-only parameters without safe_gate, and safe_gate with a linear gate unprotected against regressions. Focused rejection tests would prevent future changes from passing invalid combinations to the kernel.
Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!
1 participant
Description
Adds support for KDA attention via the fused kernel in
cudnn-frontend.Type of change
Changes
cudnn-frontend.Checklist: