Repository navigation
kda in torchtitan (training) - #4351
Merged
Merged
Conversation
[ghstack-poisoned]
liangel-02
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 27, 2026 17:16
meta-cla
Bot
added
the
CLA Signed
shuhuayu
approved these changes
Aug 27, 2026
tianyu-l
approved these changes
Aug 27, 2026
wwwjn
reviewed
Aug 28, 2026
| ngpu: int = 4 | ||
| disabled: bool = False | ||
| skip_rocm_test: bool = False | ||
| required_cuda_capabilities: Sequence[tuple[int, int]] = () |
Contributor
There was a problem hiding this comment.
We don't have blackwell for titan CI, so Kimi-K3 tests are going to be skipped anyway
Contributor
JavaZeroo
added a commit
to JavaZeroo/torchtitan
that referenced
this pull request
Aug 29, 2026
KDA already takes document offsets after pytorch#4351, but Kimi K3 still rejected MMSamplePackingConfig and handed its KDA layers a hard-coded None, so the offsets never reached them. get_attention_masks now returns a mask dict, mirroring Qwen3.5's hybrid layout: MLA layers read "quadratic_attention" (a BlockMask under flex, shared offsets under varlen) and KDA layers read "kda". The block routes each layer to its own entry, and the packing guard is gone. The offsets test uses Attention Gym's eager reference implementation so it runs on any CUDA device -- the fused kernel it ships alongside is compiled for sm_100 only, which is why the kernel test beside it is skipped here. What the test covers is the segmentation the offsets encode, which is implementation independent: a packed row reproduces the per-document runs bitwise, and the same row without offsets does not.
1 task
JavaZeroo
added a commit
to JavaZeroo/torchtitan
that referenced
this pull request
Sep 3, 2026
KDA already takes document offsets after pytorch#4351, but Kimi K3 still rejected MMSamplePackingConfig and handed its KDA layers a hard-coded None, so the offsets never reached them. get_attention_masks now returns a mask dict, mirroring Qwen3.5's hybrid layout: MLA layers read "quadratic_attention" (a BlockMask under flex, shared offsets under varlen) and KDA layers read "kda". The block routes each layer to its own entry, and the packing guard is gone. The offsets test uses Attention Gym's eager reference implementation so it runs on any CUDA device -- the fused kernel it ships alongside is compiled for sm_100 only, which is why the kernel test beside it is skipped here. What the test covers is the segmentation the offsets encode, which is implementation independent: a packed row reproduces the per-document runs bitwise, and the same row without offsets does not.
4 of 18 tasks
JavaZeroo
added a commit
to JavaZeroo/torchtitan
that referenced
this pull request
Sep 24, 2026
KDAKernel only accepted SM100/SM103. That gate matched attn-gym 0.0.5, pinned when pytorch#4351 moved KDA onto Attention Gym. Since 0.0.7, chunk_kda runs CuTe kernels on SM100/SM103 and Triton kernels on other NVIDIA GPUs, and the pinned 0.0.10 keeps that dispatch. The remaining hardware floor is the fused bounded gate, which needs TMA (SM90+). - Relax the KDAKernel gate and the two KDA GPU test skips to >= (9, 0). - Document the hardware requirement in the Kimi K3 README.
JavaZeroo
added a commit
to JavaZeroo/torchtitan
that referenced
this pull request
Sep 24, 2026
KDAKernel only accepted SM100/SM103. That gate matched attn-gym 0.0.5, pinned when pytorch#4351 moved KDA onto Attention Gym. Since 0.0.7, chunk_kda runs CuTe kernels on SM100/SM103 and Triton kernels on other NVIDIA GPUs, and the pinned 0.0.13 keeps that dispatch. The remaining hardware floor is the fused bounded gate, which needs TMA (SM90+). - Relax the KDAKernel gate and the two KDA GPU test skips to >= (9, 0). - Document the hardware requirement in the Kimi K3 README.
tianyu-l
pushed a commit
that referenced
this pull request
Sep 25, 2026
Currently, `KDAKernel` only accept SM100/SM103, so Kimi K3 and its KDA tests only run on B200-class GPUs. This gate is stale: - #4351 added it with the `attn-gym==0.0.5` pin, when the KDA kernels were Blackwell-only. - Since attn-gym 0.0.7, `chunk_kda` uses CuTe kernels on SM100/SM103 and Triton kernels on other NVIDIA GPUs. #4374 was closed for this reason. K3's real hardware floor is the fused bounded gate (`bound_gate(impl="fused")`), which needs TMA, i.e. CUDA capability 9.0+. ### Changes - `KDAKernel`: require CUDA capability >= 9.0 instead of exactly SM100/SM103. - `test_kda_attention.py`, `test_kimi_k3.py`: relax the KDA skip conditions to match. - Kimi K3 README: document the hardware requirement. SM100/SM103 behavior is unchanged
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
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.
(remaking bc i accideltny merged #4164 into my ghstack instead of main)
this PR adds attention module for KDA (relying on attention gym KDA kernels and user facing APIs)
correctness
tested with c4 dataset:
torchtitan attention gym vs hugging face (natural routing)
^i compared hugigng face and torchtitan layer by layer using identical weights and inputs, first divergence appears at block 2 where expert choice agreement was 94.5% and max router score diff was 0.0065. with nearly tied experts, small numerical differences led to different experts being in the top k which led to model executing different expert weights.
if we force expert routing to be the same, we see more alignment. i did this by first running hugging face and recording the top k experts for every token at every layer and then replacing torchtitan with those choices.
torchtitan attention gym vs hugging face (forced expert routing)
torchtitan fla vs hugging face
performance
^attention gym is 1.38x and uses 26% less memory
attention gym trace:
fla trace:
Stack from ghstack (oldest at bottom):