Repository navigation
[Test] aiter#6912 rebased on main (cu_seqlens_cpu for KDA) - #14
Closed
amd-sriram wants to merge 1 commit into
Closed
amd-sriram wants to merge 1 commit into
amd-sriram wants to merge 1 commit into
Conversation
chunk_kimi_delta_attn read cu_seqlens back to the host twice per new
tensor on the FlashKDA route (prepare_chunk_indices and _seq_bounds),
and on a descriptor-cache miss uploaded the segment tables with two
blocking copies. The first KDA layer of every serving step therefore
waited for all work queued before it.
- Accept fla's cu_seqlens_cpu on chunk_kimi_delta_attn and thread it
to chunk_delta_attn_fwd and flash_kda_fwd. With it, cu_seqlens is
never read back; without it, it is read once per tensor.
- Build the chunk indices and the segment descriptors from the host
bounds and upload them with one non-blocking copy from pinned
memory, cached on the bounds' contents. A hit from another stream
waits on the uploading stream and records the buffer on its own. A
miss during graph capture raises; a hit keeps its buffer for the
life of the process, so cache eviction cannot free it under a graph.
- prepare_chunk_indices now numbers chunks by their sequence's
position in cu_seqlens. The old cumsum-of-chunk-starts formula
skipped empty sequences, so an empty sequence before a non-empty
one sent the later sequences' chunks to the wrong bounds. The fix
covers both the FlashKDA route and the default pipeline.
Tests: a call with host bounds makes no synchronizing call under
torch.cuda.set_sync_debug_mode("error") on a metadata-cache miss, the
same call with device bounds does, and both give bit-identical output
and state; prepare_chunk_indices numbers empty sequences like a
repeat_interleave reference.
Signed-off-by: Rony Leppänen <rony.leppanen@amd.com>
Co-authored-by: Claude <noreply@anthropic.com>
|
Important
This repository does not receive automatic reviews because it has fewer than 10 stars.
 |
Owner
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
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.
Test vehicle only. ROCm#6912 (eppaneamd, draft) no longer applies on main after ROCm#6609 (gfx1250 FlashKDA). This is its first commit cherry-picked on main with the conflicts resolved by keeping both sides:
.contiguous()line, then [Triton/Gluon] KDA: accept cu_seqlens_cpu, upload varlen metadata async ROCm/aiter#6912'sprepare_chunk_indices(..., cu_seqlens_cpu).The bench_flash_kda.py commit is left out (benchmark only). Used so DockerMaker images on current AITER main still carry ROCm#6912. Close once ROCm#6912 is rebased upstream.