Repository navigation
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>
--varlen LENS times one B=1 varlen call, giving every call a new cu_seqlens (pinned host tensor, non-blocking copy) from 128 rotating bound sets, so the bounds and metadata caches miss as in a serving step. --seqlens host also passes cu_seqlens_cpu; the dense sweep is unchanged. Signed-off-by: Rony Leppänen <rony.leppanen@amd.com> Co-authored-by: Claude <noreply@anthropic.com>
Contributor
🏷️ CI GuideRuns automatically on every PR:
Extended tests (opt-in via labels):
One backend per PR: PR title tags & labels: |
5 of 6 tasks
…ens-cpu Conflicts with #6609 (gfx1250 FlashKDA, ffddbaf), resolved by keeping both sides: - chunk_fwd.py: main packs q, k, v and g after the FlashKDA return; the default pipeline then builds chunk_indices from cu_seqlens_cpu. - chunk_delta_attn.py: cu_seqlens_cpu is accepted (no longer listed as rejected); main's contiguity note, with the gfx1250 exception, is kept. - test_flash_kda.py: both test blocks kept. #6609's test_positional_call_keeps_skipped_storage lists every flash_kda_fwd parameter, so it gains cu_seqlens_cpu. - bench_flash_kda.py: --varlen's column is named "bounds", since --k3-tp8's "seqlens" column holds the lengths. --varlen rejects --k3-tp8 (its inputs carry their own cu_seqlens) and --impl chunk_kda (it would time flash_kda_fwd). The gfx1250 K1 and K2 read chunk_indices and the segment descriptors as the gfx950 kernels do, so they take the int32 metadata from _build_metadata unchanged. Co-authored-by: Claude <noreply@anthropic.com> Signed-off-by: Rony Leppänen <rony.leppanen@amd.com>
waqahmed-amd-fi
added a commit
that referenced
this pull request
Oct 28, 2026
6 tasks done
NOT READY FOR HUMAN REVIEWPR Self-ReviewRevision Waiting for the PR author. The readiness report will appear after the self-review answers are complete for this revision. Questions to answer in the PR description
Edit the Missing or invalid answers
Shadow mode reports what would block but keeps the status successful. Set |
Contributor
There was a problem hiding this comment.
🟡 Changes recommended
The staging tensor is not explicitly allocated on CPU, causing pin_memory() to fail on metadata-cache misses when CUDA is the default device.
3 open issues
What changed in this PR
This PR reduces host synchronization in KDA variable-length attention by accepting caller-provided CPU sequence bounds and uploading cached metadata asynchronously. It updates host-side preparation without changing kernel bodies.
Changes:
- Add optional
cu_seqlens_cpusupport to FlashKDA and the default KDA pipeline. - Combine FlashKDA metadata into one aligned upload, manage stream and graph lifetimes, and correct chunk numbering after empty sequences.
- Add synchronization and numbering tests, plus a variable-length benchmark with rotating bounds.
| File | Description |
|---|---|
op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py |
Test host-bound synchronization avoidance and empty-sequence numbering. |
op_tests/op_benchmarks/triton/bench_flash_kda.py |
Add variable-length benchmarks with host or device bounds. |
aiter/ops/triton/kimi_delta_attn/chunk_delta_attn.py |
Expose, validate, and forward optional CPU bounds. |
aiter/ops/triton/_triton_kernels/kimi_delta_attn/utils/index.py |
Share bounds reads and cache correctly numbered chunk indices. |
aiter/ops/triton/_triton_kernels/kimi_delta_attn/flash_kda.py |
Pack chunk and segment metadata into one aligned upload. |
aiter/ops/triton/_triton_kernels/kimi_delta_attn/chunk_fwd.py |
Forward CPU bounds through both execution paths. |
aiter/ops/triton/_triton_kernels/kimi_delta_attn/chunk_delta_attn_utils.py |
Add asynchronous uploads, stream waits, and graph-lifetime retention. |
Review effort: Balanced
| "capturing so it is cached." | ||
| ) | ||
| stream = torch.cuda.current_stream(device) | ||
| host = torch.tensor(values, dtype=dtype).pin_memory() |
| return | ||
| current = torch.cuda.current_stream(stream.device) | ||
| if current != stream: | ||
| current.wait_stream(stream) |
Comment on lines
+150
to
+152
| if torch.cuda.is_current_stream_capturing(): | ||
| _captured_uploads[buf.data_ptr()] = buf | ||
| return |
This branch has not been deployed
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.
Motivation
A serving engine calls
chunk_kimi_delta_attnonce per KDA layer per step, with a newcu_seqlenseach step. On the FlashKDA route the first KDA layer of a step readscu_seqlensback to the host twice (prepare_chunk_indicesand_seq_bounds) and, on a descriptor-cache miss, uploads the segment tables with two blocking copies. Each of these waits for all GPU work queued before that layer. Serving engines already hold the same bounds on the host.Separately,
prepare_chunk_indicesnumbers chunks wrongly when an empty sequence comes before a non-empty one: the old formula counted only sequences with chunks, so the later sequences' chunks are assigned to earlier sequences' bounds.Changes
chunk_kimi_delta_attn(..., cu_seqlens_cpu=None), with the name and meaning of fla'schunk_kdaargument: a CPU tensor shaped likecu_seqlens, trusted to hold the same values. With it,cu_seqlensis never read back; without it, it is read back once per tensor instead of twice.flash_kda.py:_build_segmentsbecomes_build_metadata, which also builds K1's chunk table. The three tables share one int32 buffer, uploaded with one non-blocking copy from pinned memory. Each block starts 16-byte aligned, becausefast_launchspecializes on pointer alignment. K1's chunk table is therefore int32 whatever the dtype ofcu_seqlens; the Triton and Gluon K1 kernels both convert it on load.prepare_chunk_indices(cu_seqlens, chunk_size, cu_seqlens_cpu=None)is built from the bounds and cached on them, and numbers chunks by their sequence's position incu_seqlens, as fla does. Dtype, device and shape are unchanged.RuntimeError.None.op_tests/op_benchmarks/triton/bench_flash_kda.py(second commit):--varlen LENStimes oneB=1varlen call over those lengths; each call gets new bounds from a rotating set of 128, so the bounds and metadata caches miss as they do in serving.--seqlens hostalso passescu_seqlens_cpu. The dense sweep is unchanged.Scope
Model-independent host code. It reaches
aiter.ops.triton.kimi_delta_attn.chunk_kimi_delta_attn(FlashKDA and the default pipeline) andaiter.ops.triton.attention.chunk_kdaon archs other than gfx950 and gfx1250, which falls back to it without host bounds and gets the single read-back, the async upload and the numbering fix. Gated DeltaNet's ownprepare_chunk_indicesis not changed. KDA layers appear in Kimi-Linear, Kimi-K3 and GLM-5.3-Flash, the measured model.Tests
Added to
op_tests/triton_tests/kimi_delta_attn/test_flash_kda.py:test_cu_seqlens_cpu_issues_no_sync: a paged FlashKDA call with host bounds makes no synchronizing call on a metadata-cache miss undertorch.cuda.set_sync_debug_mode("error"), the same call with device bounds raises, and the two give identical output and state.test_prepare_chunk_indices_numbers_empty_sequences: chunk numbering matches arepeat_interleavereference when empty sequences come before non-empty ones.Run on MI355X (gfx950, one GPU; torch 2.13.0, Triton 3.8.0), with this branch's five changed source files overlaid on AITER 0.1.26:
test_cu_seqlens_cpu_issues_no_syncfails withTypeError(nocu_seqlens_cpuargument) andtest_prepare_chunk_indices_numbers_empty_sequencesfails at its value assertion.op_tests/triton_tests/kimi_delta_attn/: 137 passed, 1 skipped on this branch; 135 passed, 2 failed (the two new tests), 1 skipped without the change. The skip needs a second GPU.Black and Ruff are clean on the changed files.
Performance
The gap this targets was seen in a profile taken before this change:
amd/GLM-5.3-Flash-Quark-MXFP4(34 KDA layers).prepare_chunk_indicesmakes one more blocking copy than this PR's base.chunk_kimi_delta_attn(FlashKDA, pagedstate_cache), now [ROCm][Perf] Add AITER FlashKDA prefill backend for GLM KDA vllm-project/vllm#63568.In each step with prefill, the GPU idles 1.27 ms just before the first KDA layer's FlashKDA kernels: 1.10 % of profiled time. Later KDA layers show no such gap. The profile does not separate host-side planning from a stream drain.
A standalone step emulation on this branch (MI355X; 34 KDA layers, 32 heads, one new
cu_seqlensper step for 12 decodes and prefills of 1100, 1300 and 900 tokens; one 0.1 ms GEMM queued before each step) shows the copy is gone: with device bounds every traced step has one synchronous device-to-host copy before the first KDA layer, and withcu_seqlens_cputhe trace has none. It does not measure a gain: too little GPU work was queued ahead of the step for the copy to wait on.Varlen microbenchmark (
--varlen): one call per step above, 12 decodes plus prefills of 1100, 1300 and 900 tokens (3312 tokens), with new bounds on each call:1e2d557d78), device bounds; the base has nocu_seqlens_cpu, so--seqlens hostexits with an error there.cu_seqlens_cpu).python op_tests/op_benchmarks/triton/bench_flash_kda.py --shape 1 3312 32 128 128 \ --varlen 1,1,1,1,1,1,1,1,1,1,1,1,1100,1300,900 --warmup-ms 2000 --seqlens device # or hostcu_seqlens_cpu)As a regression check, the script's default sweep of 11 fixed-length shapes, one run per arm: 10 of 11 shapes are within ±0.8 % of the base, and 1×256 (H 12) is +2.0 µs (+6.0 %) on a 33 µs call.
The end-to-end impact has not been measured.
Related PRs
prepare_chunk_indices(cu_seqlens, BT); that call keeps working unchanged.cu_seqlens_cpu.This change was developed with AI assistance (Claude); I reviewed every line and ran the tests listed.