Replace the KDA inter-chunk recurrence with a warp-specialized kernel - #344
Conversation
a120ff9 to
3f3f755
Compare
meta-cla
Bot
added
the
CLA Signed
670dac1 to
2b51897
Compare
![chatgpt-codex-connector[bot]](https://avatars.githubusercontent.com/in/1144995?s=60&v=4){kind=small}
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: a120ff9354
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
3f3f755 to
72e6f14
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 72e6f142a8
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
|
|
||
| def _delta_h_fake_common(k: torch.Tensor, u: torch.Tensor, capacity: int): | ||
| h = k.new_empty(k.shape[0], capacity, k.shape[2], k.shape[3], u.shape[-1]) | ||
| return h, u.new_empty(u.shape) |
There was a problem hiding this comment.
When u has an accepted TMA-compatible, non-overlapping dense layout with nonstandard strides (for example, a permuted token/head layout on a complete-chunk call), the real launcher creates v_new with torch.empty_like(u) and therefore preserves those strides, while this fake path reports a contiguous tensor via new_empty. Under torch.compile, Inductor can consequently make incorrect layout assumptions about v_new, and opcheck's fake-tensor check will detect the metadata mismatch; construct the fake result with empty_like as well.
AGENTS.md reference: AGENTS.md:L83-L84
Useful? React with 👍 / 👎.
72e6f14 to
3941bd3
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 3941bd3b15
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| b_w = tl.load( | ||
| w + ptr_offset((o_t[:, None], o_k1[None, :]), (H * K, 1)), | ||
| mask=m_w, | ||
| w + ptr_offset((o_t[:, None], i_h, o_k[None, :]), (H * K, K, 1)), |
There was a problem hiding this comment.
When a ragged sequence has a partial final chunk and k, w, or u uses a TMA-compatible nonstandard dense layout (for example, a head-major permutation with a contiguous last dimension), the wrapper accepts it via can_use_tma, but this peeled path computes addresses using hard-coded contiguous strides; the full-chunk descriptor path honors the actual strides, so only the tail silently reads and writes incorrect elements. Pass the real strides to these pointer accesses or reject such layouts before launch.
Useful? React with 👍 / 👎.
3941bd3 to
459779b
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: 459779b5bf
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| _DELTA_H_ARGS = ( | ||
| "(Tensor k, Tensor w, Tensor u, Tensor gk, Tensor? initial_state, " | ||
| "Tensor? cu_seqlens, Tensor? chunk_offsets, int capacity)" | ||
| ) |
There was a problem hiding this comment.
When chunk_gated_delta_rule_fwd_h is captured with torch.compile(dynamic=True), capacity is derived from symbolic token or sequence dimensions, but this schema declares it as a concrete int. The dispatcher therefore has to concretize/specialize that value (and may reject a symbolic value on stricter builds), so reusing the compiled recurrence across different token or packed-sequence counts cannot retain one dynamic graph. Declare this argument as SymInt capacity so output shapes remain symbolic through the fake implementation and operator boundary.
AGENTS.md reference: AGENTS.md:L83-L84
Useful? React with 👍 / 👎.
459779b to
2c969db
Compare
The blockdim64 recurrence kernel was latency-bound: 1 CTA/SM, ~77% no-eligible scheduler cycles, and a config sweep (BV up to 128, 4-32 warps, 2-4 stages) bottomed out at 248us fixed / 268us mixed at the h64/T=8192 contract shapes. The replacement holds the full [K, BV] state in one accumulator (two dots per chunk instead of four) and walks full chunks with a warp-specialized tensor-descriptor loop, so next-chunk loads overlap the serially dependent state MMAs. Partial tail chunks use a peeled masked-pointer epilogue because a descriptor load would cross sequence boundaries in the token-major packed layout, and a small Triton kernel premultiplies exp2 of each chunk's last-row cumulative gate. Output is bitwise-identical to blockdim64 (validated on dense and ragged shapes with tails and initial state before deleting it). Host-side tensor descriptors do not survive dynamo/inductor mutation analysis, so the launch sits behind a compiler-opaque torch.library op pair (attn_gym::kda_delta_h / kda_delta_h_with_state, following the kda_chunk_fwd fixed-arity precedent) instead of being traced directly; fullgraph compilation runs the same kernel through the op. With the op boundary in place, blockdim64 had no remaining role and is deleted. The wrapper contract narrows to what production uses and validates loudly: BF16, batch 1 (dense batches are packed upstream), K=V=128, 64-token chunks, and CUDA capability 10.0+; oversized tensors promote the pointer paths to int64 in-kernel, and empty inputs short-circuit on the host. The blockdim64 kernel tests are ported to fp64-golden-reference tests against the public wrapper, including the partial-tail case. GB300, warm CUDA events, h64/T=8192 contract shapes; bitwise-identical outputs: | shape | blockdim64 | warp-specialized | |---|---:|---:| | fixed [8192] | 247.6us | 135.3us | | mixed [1300,547,2048,963,271,3063] | 293.9us | 202.2us | ```bash pytest test/test_kda*.py python agent_space/bench_three_way.py # six-shape contract, numerics vs FlashInfer <= 0.001 ``` stack-info: PR: #344, branch: drisspg/stack/83
2c969db to
cc66705
Compare
There was a problem hiding this comment.
💡 Codex Review
Here are some automated review suggestions for this pull request.
Reviewed commit: cc667051f1
ℹ️ About Codex in GitHub
Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you
- Open a pull request for review
- Mark a draft as ready
- Comment "@codex review".
If Codex has suggestions, it will comment; otherwise it will react with 👍.
Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".
| # TF32 mantissa error compounds through the sequential state updates, so | ||
| # this is a path-coverage check at TF32 tolerances, not a precision | ||
| # guarantee (the dots ran in TF32 before this kernel, too). | ||
| tf32 = dict(atol=1e-1, rtol=5e-2) |
There was a problem hiding this comment.
Replace this fixed atol=0.1, rtol=0.05 budget with a data-derived error bound: for this FP32/TF32 path, a recurrence regression producing errors as large as 0.1 absolute or 5% relative can still pass all three assertions, making the new path-coverage test too permissive to protect numerical correctness. Extend or use the shared low-precision-reference helper so the allowance tracks the expected TF32 error for these inputs instead of a hand-tuned constant.
AGENTS.md reference: test/AGENTS.md:L5-L13
Useful? React with 👍 / 👎.
1 participant
Stacked PRs:
Replace the KDA inter-chunk recurrence with a warp-specialized kernel
The blockdim64 recurrence kernel was latency-bound: 1 CTA/SM, ~77% no-eligible scheduler cycles,
and a config sweep (BV up to 128, 4-32 warps, 2-4 stages) bottomed out at 248us fixed / 268us mixed
at the h64/T=8192 contract shapes. The replacement holds the full [K, BV] state in one accumulator
(two dots per chunk instead of four) and walks full chunks with a warp-specialized tensor-descriptor
loop, so next-chunk loads overlap the serially dependent state MMAs. Partial tail chunks use a peeled
masked-pointer epilogue because a descriptor load would cross sequence boundaries in the token-major
packed layout, and a small Triton kernel premultiplies exp2 of each chunk's last-row cumulative gate.
Output is bitwise-identical to blockdim64 (validated on dense and ragged shapes with tails and
initial state before deleting it).
Host-side tensor descriptors do not survive dynamo/inductor mutation analysis, so the launch sits
behind a compiler-opaque torch.library op pair (attn_gym::kda_delta_h / kda_delta_h_with_state,
following the kda_chunk_fwd fixed-arity precedent) instead of being traced directly; fullgraph
compilation runs the same kernel through the op. With the op boundary in place, blockdim64 had no
remaining role and is deleted. The wrapper contract narrows to what production uses and validates
loudly: BF16, batch 1 (dense batches are packed upstream), K=V=128, 64-token chunks, and CUDA
capability 10.0+; oversized tensors promote the pointer paths to int64 in-kernel, and empty inputs
short-circuit on the host. The blockdim64 kernel tests are ported to fp64-golden-reference tests
against the public wrapper, including the partial-tail case.
GB300, warm CUDA events, h64/T=8192 contract shapes; bitwise-identical outputs: