Skip to content

refactor(cake_causal_conv1d): deterministic names for single-rank training package - #8074

Closed
yyihuang wants to merge 6 commits into
flashinfer-ai:mainfrom
yyihuang:conv1d-det-train-names
Closed

refactor(cake_causal_conv1d): deterministic names for single-rank training package#8074
yyihuang wants to merge 6 commits into
flashinfer-ai:mainfrom
yyihuang:conv1d-det-train-names

Conversation

@yyihuang

@yyihuang yyihuang commented Oct 31, 2026

Copy link
Copy Markdown
Collaborator

This draft delivers Cake's single-rank packed width-4 depthwise causal conv1d + SiLU training under deterministic names. It covers forward, backward with an FP32 or BF16 weight gradient, and autograd. Only names change. Every generated program keeps its digest, and its kernel source changes only in the name tokens.

The entry module is now flashinfer.cake_causal_conv1d_det_train, and its package is flashinfer.cake_causal_conv1d_det_train_backend. They provide cake_causal_conv1d_det, CakeCausalConv1dDetFunction, cake_causal_conv1d_det_train_forward, cake_causal_conv1d_det_train_backward, their _workspace / _workspace_bytes helpers, and is_cake_causal_conv1d_det_train_supported. The previous names still point to the same objects: cake_causal_conv1d, CakeCausalConv1dFunction, cake_causal_conv1d_train_forward, cake_causal_conv1d_train_backward, the workspace helpers, and is_cake_causal_conv1d_train_supported. The previous module paths flashinfer.cake_causal_conv1d_train and flashinfer.cake_causal_conv1d_train_backend still import, and they resolve to the new modules. The kernel libraries for the 29 single-rank programs are renamed from cake_causal_conv1d_train_<digest> to cake_causal_conv1d_det_train_<digest>.

The context-parallel variant is renamed in the stacked PR #8055, and the supported targets are sm_100a, sm_103a and sm_107a.

Speedup tables follow after measurement.

🤖 Generated with Claude Code

@coderabbitai

coderabbitai Bot commented Oct 31, 2026 •
edited
Loading

Copy link
Copy Markdown
Contributor

Important

Draft PR not reviewed

Draft PRs are not automatically reviewed by default.

  • Trigger a manual review

To automatically review draft PRs, update your CodeRabbit configuration:

reviews:
  auto_review:
    drafts: true
  • Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

Share

@flashinfer-bot flashinfer-bot added op: misc norm, activation, sampling, RoPE, quantization, etc. op: linear attention KDA, mamba, GDN, etc. review filtering. labels Oct 31, 2026
@yyihuang yyihuang mentioned this pull request Oct 31, 2026
Merged
@yyihuang
yyihuang force-pushed the conv1d-det-train-names branch 2 times, most recently from 1f83919 to 0391b23 Compare October 31, 2026 00:40
averyhNV and others added 4 commits October 30, 2026 19:44
@averyhNV @claude
chore(causal_conv1d): move the single-rank training package to the de…
b05d822 
…t names

The package flashinfer.cake_causal_conv1d_train_backend moves to
flashinfer.cake_causal_conv1d_det_train_backend. Its single-rank programs
take the cake_causal_conv1d_det_train prefix: the module registry, the
kernel and binding units and their build.json entries are renamed, and
build.json records the digests of the renamed units. Every program keeps
its twenty hex digits, kernel code and launch rule. The KCP programs keep
their names and bytes.

Co-Authored-By: Claude <noreply@anthropic.com>
@averyhNV @claude
refactor(causal_conv1d): det names for the single-rank training entry…
46793b4 
… points

The packed width-4 causal convolution + SiLU training module is now
flashinfer.cake_causal_conv1d_det_train, and its backend package is
flashinfer.cake_causal_conv1d_det_train_backend. The single-rank entry points
are cake_causal_conv1d_det, cake_causal_conv1d_det_train_{forward,backward},
cake_causal_conv1d_det_train_{forward,backward}_workspace[_bytes],
is_cake_causal_conv1d_det_train_supported and CakeCausalConv1dDetFunction.
The facade, the backend package and the flashinfer top level also bind the
previous names to the same objects. These names are not in __all__.

The generated sources, build.json and the two registries move without change.
The tests, the benchmark, the API docs, the package data, the build manifest
check and the pre-commit exclusion use the new paths.

Co-Authored-By: Claude <noreply@anthropic.com>
@averyhNV @claude
feat(causal_conv1d): compatibility imports for the previous module paths
23f4983 
flashinfer.cake_causal_conv1d_train is the det facade's module object.
flashinfer.cake_causal_conv1d_train_backend resolves its attributes on the det
backend package at first access. Its cake_build, cake_jit, kcp and kcp_jit
submodules are the det package's module objects, so the registries and the
workspace caches exist once. The compatibility package holds no sources and
no manifest, and importing it does not import the KCP host module.

tests/norm/test_cake_causal_conv1d_det_train_aliases.py checks the identity of
every compatibility name in its own namespace, the module identity of the
compatibility imports, and that no import form loads kcp or kcp_jit under
either path.

Co-Authored-By: Claude <noreply@anthropic.com>
@averyhNV @claude
fix(causal_conv1d): previous module paths resolve as flashinfer attri…
c2db120 
…butes

After a plain `import flashinfer`, `flashinfer.cake_causal_conv1d_train` and
`flashinfer.cake_causal_conv1d_train_backend` resolve on first attribute access
through the package `__getattr__`, as the facade and its backend package did
before the det names. Neither access loads the KCP host module. The alias test
checks this attribute form in a fresh interpreter.

Co-Authored-By: Claude <noreply@anthropic.com>
@yyihuang
yyihuang force-pushed the conv1d-det-train-names branch from 0391b23 to c2db120 Compare October 31, 2026 02:45
@yyihuang

yyihuang commented Oct 31, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_causal_conv1d_det_train.py tests/norm/test_cake_causal_conv1d_det_train_aliases.py tests/norm/test_cake_causal_conv1d_det_train_api.py

@yyihuang

yyihuang commented Oct 31, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_causal_conv1d_det_train.py tests/norm/test_cake_causal_conv1d_det_train_aliases.py tests/norm/test_cake_causal_conv1d_det_train_api.py

yyihuang added a commit that referenced this pull request Oct 31, 2026
@yyihuang @averyhNV @claude
refactor(cake_causal_conv1d): deterministic KCP names for the context…
dd8325c 
…-parallel causal conv1d entry points (#8055)

This PR delivers Cake's context-parallel (KCP) packed width-4 depthwise
causal conv1d + SiLU with halo exchange under deterministic names. It
covers forward, backward and autograd, for one tensor or for several
tensors of one layer. Only names change, and the generated programs keep
their digests.

The new entry module `flashinfer.cake_causal_conv1d_det_train_kcp`
provides `cake_causal_conv1d_det_kcp`,
`cake_causal_conv1d_det_multi_kcp` and `RecomputeHandleDetKcp`. The
programs stay in `flashinfer.cake_causal_conv1d_det_train_backend`.
Their `kcp` names there are
`cake_causal_conv1d_det_{forward,backward}_kcp`,
`cake_causal_conv1d_det_multi_{forward,backward}_kcp`,
`CakeCausalConv1dDetFunctionKcp` and
`CakeCausalConv1dDetMultiFunctionKcp`. The previous names still point to
the same objects, including in
`flashinfer.cake_causal_conv1d_det_train`. Examples are
`cake_causal_conv1d_kcp`, `cake_causal_conv1d_kcp_multi`,
`KcpRecomputeHandle` and `CakeCausalConv1dKcpFunction`. Kernel libraries
are renamed `cake_causal_conv1d_det_train_<digest>_kcp`. All five are
used only for context parallelism. The supported targets are sm_100a,
sm_103a and sm_107a.

This PR also carries the complete single-rank deterministic rename and
supersedes #8074. The single-rank names are
`flashinfer.cake_causal_conv1d_det_train`, `cake_causal_conv1d_det`,
`CakeCausalConv1dDetFunction` and the deterministic forward/backward and
workspace entry points. All previous names and module paths remain
compatible aliases.

**Single-rank speed.** "Before" is FI main under the previous names;
"after" is this PR. S is the packed token count and C the channel count.
A ratio below 1 means this PR is faster. After/before ranges from 0.9958
to 1.0056, so speed is unchanged within run-to-run noise.

NVIDIA B200 (sm_100a), CUDA 13.4, torch 2.14.1+cu130

| S | C | pass | before (ms) | after (ms) | after/before |
|---:|---:|---|---:|---:|---:|
| 239718 | 512 | fwd | 0.08018 | 0.08020 | 1.0002 |
| 239718 | 512 | bwd | 0.14197 | 0.14194 | 0.9998 |
| 239718 | 512 | fwd+bwd | 0.22162 | 0.22159 | 0.9999 |
| 245028 | 512 | fwd | 0.08181 | 0.08179 | 0.9997 |
| 245028 | 512 | bwd | 0.14326 | 0.14326 | 1.0000 |
| 245028 | 512 | fwd+bwd | 0.22449 | 0.22449 | 1.0000 |
| 249866 | 512 | fwd | 0.08324 | 0.08326 | 1.0002 |
| 249866 | 512 | bwd | 0.14757 | 0.14752 | 0.9997 |
| 249866 | 512 | fwd+bwd | 0.23061 | 0.23074 | 1.0006 |
| 250340 | 512 | fwd | 0.08350 | 0.08351 | 1.0001 |
| 250340 | 512 | bwd | 0.14655 | 0.14653 | 0.9999 |
| 250340 | 512 | fwd+bwd | 0.22917 | 0.22920 | 1.0001 |
| 258574 | 512 | fwd | 0.08592 | 0.08590 | 0.9998 |
| 258574 | 512 | bwd | 0.14991 | 0.14990 | 0.9999 |
| 258574 | 512 | fwd+bwd | 0.23558 | 0.23560 | 1.0001 |
| 259042 | 512 | fwd | 0.08604 | 0.08604 | 0.9999 |
| 259042 | 512 | bwd | 0.15030 | 0.15032 | 1.0001 |
| 259042 | 512 | fwd+bwd | 0.23580 | 0.23580 | 1.0000 |
| 259584 | 512 | fwd | 0.08620 | 0.08620 | 1.0000 |
| 259584 | 512 | bwd | 0.15209 | 0.15206 | 0.9998 |
| 259584 | 512 | fwd+bwd | 0.23812 | 0.23812 | 1.0000 |
| 261393 | 512 | fwd | 0.08657 | 0.08659 | 1.0002 |
| 261393 | 512 | bwd | 0.14997 | 0.15000 | 1.0002 |
| 261393 | 512 | fwd+bwd | 0.23688 | 0.23697 | 1.0004 |
| 250340 | 128 | fwd | 0.02943 | 0.02942 | 0.9996 |
| 250340 | 128 | bwd | 0.05116 | 0.05121 | 1.0009 |
| 250340 | 128 | fwd+bwd | 0.08323 | 0.08351 | 1.0034 |
| 250340 | 256 | fwd | 0.04748 | 0.04748 | 1.0002 |
| 250340 | 256 | bwd | 0.08388 | 0.08386 | 0.9998 |
| 250340 | 256 | fwd+bwd | 0.13305 | 0.13380 | 1.0056 |
| 250340 | 1024 | fwd | 0.16458 | 0.16461 | 1.0002 |
| 250340 | 1024 | bwd | 0.26882 | 0.26887 | 1.0002 |
| 250340 | 1024 | fwd+bwd | 0.43340 | 0.43309 | 0.9993 |
| 250340 | 2048 | fwd | 0.31720 | 0.31786 | 1.0021 |
| 250340 | 2048 | bwd | 0.52717 | 0.52721 | 1.0001 |
| 250340 | 2048 | fwd+bwd | 0.89177 | 0.88991 | 0.9979 |
| 250340 | 4096 | fwd | 0.64286 | 0.64296 | 1.0002 |
| 250340 | 4096 | bwd | 1.09285 | 1.09101 | 0.9983 |
| 250340 | 4096 | fwd+bwd | 1.83452 | 1.82674 | 0.9958 |
| 250340 | 8192 | fwd | 1.33790 | 1.34240 | 1.0034 |
| 250340 | 8192 | bwd | 2.27470 | 2.27473 | 1.0000 |
| 250340 | 8192 | fwd+bwd | 3.75009 | 3.75094 | 1.0002 |
| 250340 | 1280 | fwd | 0.19944 | 0.19946 | 1.0001 |
| 250340 | 1280 | bwd | 0.34438 | 0.34487 | 1.0014 |
| 250340 | 1280 | fwd+bwd | 0.55374 | 0.55403 | 1.0005 |
| 262144 | 1280 | fwd | 0.20840 | 0.20840 | 1.0000 |
| 262144 | 1280 | bwd | 0.36126 | 0.36146 | 1.0006 |
| 262144 | 1280 | fwd+bwd | 0.58010 | 0.58030 | 1.0004 |
| 250340 | 1536 | fwd | 0.23050 | 0.23051 | 1.0000 |
| 250340 | 1536 | bwd | 0.40193 | 0.40197 | 1.0001 |
| 250340 | 1536 | fwd+bwd | 0.65502 | 0.65504 | 1.0000 |
| 262144 | 1536 | fwd | 0.24097 | 0.24098 | 1.0000 |
| 262144 | 1536 | bwd | 0.42038 | 0.42043 | 1.0001 |
| 262144 | 1536 | fwd+bwd | 0.68282 | 0.68423 | 1.0021 |
| 250340 | 2560 | fwd | 0.37849 | 0.37856 | 1.0002 |
| 250340 | 2560 | bwd | 0.66064 | 0.66168 | 1.0016 |
| 250340 | 2560 | fwd+bwd | 1.11339 | 1.11441 | 1.0009 |
| 250340 | 1344 | fwd | 0.21015 | 0.21014 | 0.9999 |
| 250340 | 1344 | bwd | 0.40603 | 0.40602 | 1.0000 |
| 250340 | 1344 | fwd+bwd | 0.62640 | 0.62695 | 1.0009 |
| 250340 | 1160 | fwd | 0.19931 | 0.19930 | 0.9999 |
| 250340 | 1160 | bwd | 0.36618 | 0.36648 | 1.0008 |
| 250340 | 1160 | fwd+bwd | 0.56317 | 0.56316 | 1.0000 |
| 262144 | 1160 | fwd | 0.20820 | 0.20819 | 1.0000 |
| 262144 | 1160 | bwd | 0.37148 | 0.37186 | 1.0010 |
| 262144 | 1160 | fwd+bwd | 0.58721 | 0.58780 | 1.0010 |

Each value is the mean of two independent package copies; each copy's
value is the median of 48 interleaved A/B measurements (3 fresh
processes x 8 buffer placements x 2 pairings), each measurement the
median of 3 x 64 timed calls (full call span, cold L2), all on one GPU;
after/before = after (ms) / before (ms).

**Context-parallel (KCP) speed: NVIDIA GB200 (sm_100a), 4 GPUs on one
node, CUDA 13.0 runtime (torch 2.14.1).** "Before" is the same KCP path
under the previous names on the parent branch. "After" is this PR. The
ratio is after/before, so a value below 1 means this PR is faster. Each
iteration of `benchmarks/bench_cake_kda_train_kcp.py` (torchrun, one
process per rank) times the KCP causal conv1d and then the KCP KDA step,
forward and backward, with CUDA events. The slowest rank's time counts.
The KDA programs are the same on both sides. The workload is a packed
seqlen of 65536 with 3 documents (15945, 21845 and 27746 tokens), 16
heads and 4096 channels. Each value is the median over interleaved runs
per side, in the order base, branch, branch, base. Each run's figure is
the bench's median over 3 groups of 8 timed iterations.

| CP | layout | runs per side | before (ms) | after (ms) | after/before
|
|---|---|---|---|---|---|
| 2 | `contiguous` | 4 | 4.4644 | 4.4728 | 1.0019 |
| 2 | `zigzag` | 4 | 4.5155 | 4.5170 | 1.0003 |
| 2 | `zigzag_per_document` | 4 | 4.6157 | 4.6383 | 1.0049 |
| 4 | `contiguous` | 10 | 3.3509 | 3.3424 | 0.9975 |
| 4 | `zigzag` | 4 | 3.4003 | 3.3931 | 0.9979 |
| 4 | `zigzag_per_document` | 4 | 3.4441 | 3.4374 | 0.9981 |

After/before ranges from 0.9975 to 1.0049, so speed is unchanged within
run-to-run noise.

This combined PR delivers both single-rank and KCP deterministic names,
including the KCP multi-tensor form, for sm_100a, sm_103a and sm_107a.



<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->

<!-- coderabbit-ai-generated:v1 -->

## Summary by CodeRabbit

* **New Features**
* Added deterministic causal-convolution training APIs for forward,
backward, and autograd, with optional bias and configurable gradient
precision.
* Added context-parallel APIs for single- and multi-tensor workloads,
with graph replay and checkpoint recomputation support.
* Added support for empty inputs and documented the deterministic
training and context-parallel APIs.

* **Compatibility**
* Existing training API names remain available as aliases to the
deterministic APIs, including context-parallel entry points.

<!-- end of auto-generated comment: release notes by coderabbit.ai -->

---------

Co-authored-by: Yingyi Huang <averyh@nvidia.com>
Co-authored-by: Claude <noreply@anthropic.com>
@yyihuang yyihuang mentioned this pull request Oct 31, 2026
Open
8 tasks done
@yyihuang

yyihuang commented Oct 31, 2026

Copy link
Copy Markdown
Collaborator Author

Superseded by #8055, which merged at 2026-10-31T07:27:00Z as dd8325c.

#8055 delivered this PR's single-rank deterministic rename of the causal conv1d training kernels (flashinfer.cake_causal_conv1d_det_train_backend, with the former module paths kept as aliases). Both heads of this PR, 536b921 and c2db120, are ancestors of the merged #8055 head 027ee46, so no commit is lost. Closing without merging.

@yyihuang yyihuang closed this Oct 31, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@righthandabacus righthandabacus Awaiting requested review from righthandabacus righthandabacus will be requested when the pull request is marked ready for review righthandabacus is a code owner
@dierksen dierksen Awaiting requested review from dierksen dierksen will be requested when the pull request is marked ready for review dierksen is a code owner
@bkryu bkryu Awaiting requested review from bkryu bkryu will be requested when the pull request is marked ready for review bkryu is a code owner
@yongwww yongwww Awaiting requested review from yongwww yongwww will be requested when the pull request is marked ready for review yongwww is a code owner
@aleozlx aleozlx Awaiting requested review from aleozlx aleozlx will be requested when the pull request is marked ready for review aleozlx is a code owner
@yzh119 yzh119 Awaiting requested review from yzh119 yzh119 will be requested when the pull request is marked ready for review yzh119 is a code owner
@Jingfan-Sun Jingfan-Sun Awaiting requested review from Jingfan-Sun Jingfan-Sun will be requested when the pull request is marked ready for review Jingfan-Sun is a code owner
@YangXu1990uiuc YangXu1990uiuc Awaiting requested review from YangXu1990uiuc YangXu1990uiuc will be requested when the pull request is marked ready for review YangXu1990uiuc is a code owner
@Anerudhan Anerudhan Awaiting requested review from Anerudhan Anerudhan will be requested when the pull request is marked ready for review Anerudhan is a code owner
@kahyunnam kahyunnam Awaiting requested review from kahyunnam kahyunnam will be requested when the pull request is marked ready for review kahyunnam is a code owner
@flashinfer-bot flashinfer-bot Awaiting requested review from flashinfer-bot flashinfer-bot will be requested when the pull request is marked ready for review flashinfer-bot is a code owner

Assignees

No one assigned

Labels

op: linear attention KDA, mamba, GDN, etc. review filtering. op: misc norm, activation, sampling, RoPE, quantization, etc. run-ci

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@yyihuang @flashinfer-bot @averyhNV