[kimi k3] Support sample packing in KDA - #4347
Open
JavaZeroo wants to merge 5 commits into
Open
Conversation
JavaZeroo
requested review from
fegin,
tianyu-l,
wconstab and
wwwjn
as code owners
August 27, 2026 05:31
meta-cla
Bot
added
the
CLA Signed
JavaZeroo
marked this pull request as draft
August 27, 2026 05:33
tianyu-l
requested changes
Aug 27, 2026
JavaZeroo
force-pushed
the
kimi-k3-sample-packing
branch
from
August 27, 2026 06:40
6316daf to
bf909e7
Compare
Document boundaries were inferred from positions alone, but packers pad a row with position 0 on every padded token, so create_varlen_metadata_for_document read that padding as a run of one-token documents. Qwen3.5 worked around it in its own get_attention_masks with a heuristic -- a real document start is a position 0 followed by a position 1 -- which silently merges a genuine one-token document into its predecessor and only ever recovered information the packer already had. Have the dataloader state what is padding instead of having each model re-derive it: - Both packers keep Grain's input_ids_segment_ids as a padding_mask (0 marks the padding Grain appends), and both collators extend it with the batch tail they append. - preprocess_inputs pops the mask out of the batch and passes it to get_attention_masks. It describes the batch rather than the model input, so it never becomes a forward kwarg. - create_varlen_metadata_for_document takes padding_mask and collapses each contiguous padding run into one segment. Without a mask its behavior is unchanged. - Qwen3.5 drops its heuristic, including the 'single document keeps the plain kernels' short-circuit: every batch is packed and padded, so the case it guarded does not arise. This also fixes the text path, where TextCollator pads with position 0 too: every varlen-backend text model was reading a padded tail as one document per pad token. Ten deterministic steps of qwen35_debugmodel (--debug.seed=42 --debug.deterministic) reproduce loss and grad_norm exactly.
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.
FLA memoizes its varlen index helpers (prepare_chunk_indices and friends) on argument identity, with a four-entry queue. InnerGatedDeltaNet received the offsets as a host tuple and rebuilt a CPU tensor from it in every layer, so every call missed: the helpers were recomputed from scratch, and each new entry evicted the ones the rest of the kernel shares. create_varlen_metadata_for_document already reads the offsets to the host to materialize cu_seq_q_host, so keep that tensor as cu_seq_q_cpu and pass it through. Every layer then hands FLA the same object and the memoized helpers stay warm. The two 'requires a CPU cu_seqlens tensor' guards go with it: the mask builder always populates the field, and torch.where already rejects a malformed one. Profiled step 10 of qwen35_debugmodel, real training: CausalConv1dFunction 35.8 ms -> 16.3 ms ChunkGatedDeltaRule 26.4 ms -> 9.6 ms cudaStreamSynchronize 35 -> 25 calls aten::arange 124 -> 58 calls aten::narrow 142 -> 54 calls total cpu_op time 660.4 ms -> 399.0 ms Conv call count is unchanged at 72, so nothing is skipped -- the removed work is the index-table rebuild. Ten deterministic steps of qwen35_debugmodel reproduce loss and grad_norm exactly.
JavaZeroo
force-pushed
the
kimi-k3-sample-packing
branch
from
August 29, 2026 01:47
d12ac5f to
b81093f
Compare
Contributor
Author
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.
KimiDeltaAttentionnow takes document offsets as VarlenMetadata. With offsets, the Q/K/V causal convolutions run through Attention Gym's varlen causal_conv1d and the recurrence passes cu_seqlens to its chunk_kda, so both the convolution window and the recurrent state reset at document boundaries.get_attention_masksreturns 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 dataloader now emits a
padding_maskalongsidepositions: the text and multimodal packers derive it from Grain'sinput_ids_segment_ids, and their collators extend it with the batch tail they append.VarlenMetadataalso carriescu_seq_q_cpu, the host copy of the offsets that FLA's varlen kernels take, so Qwen3.5 passes one shared object per batch instead of rebuilding it in every layer.Performance
The varlen host offsets are built once per batch and carried in
VarlenMetadatainstead of being rebuilt inside every layer, which cuts both the CPU ops and the d2h per training step. Qwen3.5 is the model this shows up on: Kimi K3's KDA moved from FLA to Attention Gym in #4351 and no longer rebuilds them per layer.Qwen3.5
FLA memoizes its varlen index tables on argument identity, so rebuilding the host offsets inside each layer meant the first conv of every layer missed that cache and paid a full rebuild plus a d2h.
qwen35_debugmodelhas 12 GDN layers and calls the conv three times per layer, which is exactly what the per-call durations show:CausalConv1dFunctionIdentical in all three runs. The median call barely moves, so the win is entirely in the calls that used to rebuild.
qwen35_debugmodel, single GPU:CausalConv1dFunction(x36)ChunkGatedDeltaRuleFunction(x12)aten::repeat_interleave(x64 -> x31)aten::arange(x126 -> x58)aten::narrow(x141 -> x54)cpu_opTen deterministic steps of
qwen35_debugmodelreproduce upstream's loss and grad_norm exactly.Verification
Offsets isolate documents. Two documents (lengths 96 and 160) run separately, then their concatenation run once with
cu_seqlens = (0, 96, 256). Packed output is bitwise equal to the per-document outputs (atol=0, rtol=0). Without the offsets the second document is contaminated, as expected.End to end on B200 (sm_100,
kimi_k3_debugmodel, 10 steps,--debug.seed 42 --debug.deterministic). The 256-token debug row holds one cc12m document, so the row was widened to 2048 to actually exercise packing:Loss 12.50 -> 3.26 over 10 steps (control without the packer: 12.45 -> 4.50), grad_norm 15.1 -> 2.2.
Conclusion
Qwen3.5 is unchanged. Ten deterministic steps of
qwen35_debugmodelreproduce upstream's loss and grad_norm exactly.Kimi K3 does not reproduce upstream bitwise, and shouldn't: on a padded batch upstream splits the tail into one document per pad token, which is the bug this fixes.
TODO
.tolist()in the mask path looked like a per-step d2h -- profiled, and the real cost was the per-layer offset rebuild, now fixed.get_attention_masksin Kimi K3 and Qwen3.5 are near-identical and should be shared; needs a design.