Skip to content

refactor(cake_rmsnorm): mode-word names and experimental/cake backends (#8928) - #8910

Draft
yyihuang wants to merge 31 commits into
flashinfer-ai:mainfrom
yyihuang:fi-8066-audit-names-v1
Draft

refactor(cake_rmsnorm): mode-word names and experimental/cake backends (#8928)#8910
yyihuang wants to merge 31 commits into
flashinfer-ai:mainfrom
yyihuang:fi-8066-audit-names-v1

Conversation

@yyihuang

@yyihuang yyihuang commented Nov 6, 2026 •
edited
Loading

Copy link
Copy Markdown
Collaborator

Related issue: #8066
Naming contract: #8928 (standalone-build constraints: #8928 (comment))

What changes

This PR applies the #8928 naming contract to the training RMSNorm (Cake) packages. Every public name now carries its mode, and the name is the only place the mode appears.

  • Facades. The deterministic facade is still flashinfer.cake_rmsnorm_det_train. The non-deterministic facade is flashinfer.cake_rmsnorm_nondet_train. The old module flashinfer.cake_rmsnorm_train is removed.
  • Backends. Each backend moves as a whole directory under flashinfer/experimental/cake/. The path map is below.
  • Names. Every public name carries its mode word: det / nondet, or Det / Nondet in class names. The constants ARCHES and SUPPORTED_COMPUTE_CAPABILITIES keep their names.
  • No deterministic keyword. No entry point in either package takes deterministic. Passing it raises TypeError.
  • No aliases. The old names are removed, so code that uses them fails at import. See the migration table.
  • No cross-package fallback. Each package runs only its own kernels. The non-deterministic package never falls back to the deterministic one.
  • Behavior under torch.use_deterministic_algorithms(True):
    • The non-deterministic package alerts only before work whose result can vary from run to run. That is every backward call, plus the backward of the autograd cake_rmsnorm_nondet / CakeRMSNormNondetFunction when it runs.
    • The alert is a RuntimeError that names the det entry point. With warn_only=True, the package warns once (UserWarning) and runs the non-deterministic backward, never the deterministic one.
    • These never alert: the forward, the workspace helpers, cake_rmsnorm_nondet_train_route_key and is_cake_rmsnorm_nondet_train_supported.
    • The deterministic package never alerts.
  • Top-level flashinfer. It binds only the det names: cake_rmsnorm_det, CakeRMSNormDetFunction, cake_rmsnorm_det_train_forward, cake_rmsnorm_det_train_backward, cake_rmsnorm_det_train_backward_workspace_bytes. It also binds the two facade modules. All of these are _LAZY_EXPORTS rows. The non-deterministic package is reached through its module only.
  • No backend at import. import flashinfer imports no RMSNorm backend. Each facade loads its backend the first time it is used.

Migration

Removed Use
module flashinfer.cake_rmsnorm_train module flashinfer.cake_rmsnorm_nondet_train
cake_rmsnorm(..., deterministic=False) cake_rmsnorm_nondet(...)
CakeRMSNormFunction CakeRMSNormNondetFunction
cake_rmsnorm_train_forward cake_rmsnorm_nondet_train_forward
cake_rmsnorm_train_backward cake_rmsnorm_nondet_train_backward
cake_rmsnorm_train_backward_workspace cake_rmsnorm_nondet_train_backward_workspace
cake_rmsnorm_train_backward_workspace_bytes cake_rmsnorm_nondet_train_backward_workspace_bytes
is_cake_rmsnorm_train_supported is_cake_rmsnorm_nondet_train_supported
route_key (non-deterministic package) cake_rmsnorm_nondet_train_route_key
route_key (deterministic package) cake_rmsnorm_det_train_route_key
program_windows cake_rmsnorm_det_train_program_windows
weight_dtypes cake_rmsnorm_det_train_weight_dtypes
flashinfer.cake_rmsnorm_det_train_backend flashinfer.experimental.cake.cake_rmsnorm_det_train
flashinfer.cake_rmsnorm_train_backend flashinfer.experimental.cake.cake_rmsnorm_nondet_train

Standalone vendoring

This follows the standalone-build constraints in the issue comment linked above:

  • Each backend directory moved whole, together with:
    • its build.json, where only "package" changes;
    • cake_build.py;
    • csrc/cake_ffi_checks.h.
  • Manifest paths are package-relative.
  • No backend imports anything outside its own directory.
  • The mode is selected by the package name only.
  • The PACKAGES rows in tests/test_helpers/cake_package.py are updated to cake_rmsnorm_det_train and cake_rmsnorm_nondet_train.

Vendoring path map:

Old directory New directory Mode
flashinfer/cake_rmsnorm_det_train_backend/ flashinfer/experimental/cake/cake_rmsnorm_det_train/ det
flashinfer/cake_rmsnorm_train_backend/ flashinfer/experimental/cake/cake_rmsnorm_nondet_train/ nondet

The contents of flashinfer/cake_rmsnorm_train_backend/ depend on the version:

  • In the current release candidate, it holds the non-deterministic package.
  • In v0.7.3, it was the only package, and that package was deterministic.

Unchanged

  • Kernels, program selection and numerics are unchanged.
  • Rendered files keep their program prefixes, library stems and csrc file names.
  • The only change to rendered text is the window docstring in the non-deterministic cake_windows.py. It drops the text "(deterministic=False)".

Old names kept on purpose

A repository-wide grep still finds the old names in these places:

  • the migration table in docs/api/norm.rst;
  • the removed-name checks in tests/norm/test_cake_rmsnorm_det_train_names.py;
  • the generated program and library names in csrc, build.json and the cake_jit registries. The program prefix cake_rmsnorm_train_<hash> is kept.

Tests

Latest-head manual CI run 42232092067 at a3cc650d796efbd2067999d5f4826e362d8aa12a passed: targeted H100 tests, all six AOT build/import jobs, and the summary. GB300 validation and review remain open.

B200 (sm_100a, 148 SMs, torch 2.14.1+cu130), PR head be70eb0 merged into main 2e98955. All RMSNorm test files of both packages ran with 0 failures:

Package Tests Result
nondet test_cake_rmsnorm_nondet_train.py 851 passed, 2 skipped
nondet _all_sm 145 passed
nondet _windows 48 passed
nondet _api 46 passed
nondet _stall 109 passed
nondet _ckpt + _vendored 65 passed, 2 skipped
nondet _pairs 14 passed
nondet _dist (2 GPUs) 48 passed
det test_cake_rmsnorm_det_train.py 800 passed, 140 skipped
det _api + _names 162 passed
det _pairs 152 passed
det _training 49 passed, 1 skipped
det _stall (run alone) 66 passed
det names test (run alone) 153 passed

The same B200 at PR head 7b6e359 (merged with main 16ab4fb): names test 153 passed; JIT warm-up of both packages for sm_100a passed (159 and 147 programs). The later merges of main (heads 0eb1bb4 and 5f3f8a7) change no RMSNorm file.

GB300 (sm_103a): pending (queued).

CPU node:

  • pre-commit passes on all changed files (1333 at the final head).
  • Host-runnable RMSNorm tests: 329 passed (main: 277), with no new failures.
  • tests/cpu, tests/experimental/test_containment.py and scripts/tests/test_check_pr_api_diff.py: no new failures at the final head. In the full parallel run one timing test, tests/cpu/test_run_with_output_watchdog.py::test_watchdog_diagnoses_no_output (a 0.2 s watchdog that does not import flashinfer), failed under load; it passes when run alone.
  • scripts/check_cake_build_manifests.py: 0 findings for 44 packages.

Shared standalone-import check (#8951), with this PR merged into main: both RMSNorm packages import with flashinfer blocked. The result is 57 passed and 5 xfailed. The 5 xfails are the five allowlisted packages of other families.

@averyhNV
refactor(cake_rmsnorm): det names only; the non-deterministic package…
109e00e 
… requires deterministic=False

- flashinfer.cake_rmsnorm_det_train and its backend drop every name from
  before the det rename (CakeRMSNormFunction, cake_rmsnorm,
  cake_rmsnorm_train_forward, cake_rmsnorm_train_backward,
  cake_rmsnorm_train_backward_workspace,
  cake_rmsnorm_train_backward_workspace_bytes,
  is_cake_rmsnorm_train_supported); only det names remain.
- The det entry points take no deterministic keyword:
  cake_rmsnorm_det, cake_rmsnorm_det_train_backward and
  CakeRMSNormDetFunction no longer accept it, so passing it raises
  TypeError. torch.use_deterministic_algorithms changes nothing for them.
- flashinfer.cake_rmsnorm_train and its backend drop
  _PREVIOUS_DET_NAMES and the module __getattr__ that resolved det
  names, and the det routing and warn_only fallback; they never import
  or run the det package. Every entry point takes deterministic=True by
  default (keyword-only, positional on cake_rmsnorm and the autograd
  forward); any value other than False raises ValueError naming the det
  entry point and saying "pass deterministic=False for the
  non-deterministic path".
- Under torch.use_deterministic_algorithms(True) only the backward
  alerts (the autograd backward when it runs): RuntimeError naming the
  det entry point, or with warn_only one UserWarning, then the non-det
  backward. The forward, the workspace helpers and
  is_cake_rmsnorm_train_supported never alert.
- Docs, tests and the benchmark follow: every non-det call passes
  deterministic=False, the legacy and alias tests are replaced by
  test_cake_rmsnorm_det_train_names.py, and det arms call
  flashinfer.cake_rmsnorm_det_train explicitly.
@coderabbitai

coderabbitai Bot commented Nov 6, 2026

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

@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_train.py tests/norm/test_cake_rmsnorm_train_all_sm.py tests/norm/test_cake_rmsnorm_train_api.py tests/norm/test_cake_rmsnorm_train_ckpt.py tests/norm/test_cake_rmsnorm_train_dist.py tests/norm/test_cake_rmsnorm_train_pairs.py tests/norm/test_cake_rmsnorm_train_stall.py tests/norm/test_cake_rmsnorm_train_vendored.py tests/norm/test_cake_rmsnorm_train_windows.py

@flashinfer-bot flashinfer-bot added op: misc norm, activation, sampling, RoPE, quantization, etc. run-ci labels Nov 6, 2026
yyihuang
yyihuang commented Nov 6, 2026
View reviewed changes

@yyihuang yyihuang left a comment

Copy link
Copy Markdown
Collaborator Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Changes needed before this leaves draft:

  1. Top-level names diverge from the mHC and conv1d renames. This PR keeps flashinfer.cake_rmsnorm, flashinfer.CakeRMSNormFunction and flashinfer.cake_rmsnorm_train_* at the top level and points them at the non-deterministic package. Existing code that calls flashinfer.cake_rmsnorm(x, w) still imports and then fails at run time with ValueError. #8900 (mHC) and #8912 (conv1d) take the other route: the top level keeps only det names, the non-deterministic entries live in their own module, and old code fails at import. Please follow the same rule here: keep only the det names at the top level (cake_rmsnorm_det, CakeRMSNormDetFunction, cake_rmsnorm_det_train_*), and reach the non-deterministic package only through flashinfer.cake_rmsnorm_train. If the bare top-level names have to stay, list this run-time break under "Breaking changes".
  2. Test results are still pending. Please add the B200 and GB300 results to the description before merging.

@yyihuang yyihuang mentioned this pull request Nov 6, 2026
Draft
@averyhNV
refactor(cake_rmsnorm): the top level binds only the det names
2070071 
The flashinfer top level no longer binds cake_rmsnorm, CakeRMSNormFunction,
cake_rmsnorm_train_forward, cake_rmsnorm_train_backward and
cake_rmsnorm_train_backward_workspace_bytes. The non-deterministic package is
reached only as the module flashinfer.cake_rmsnorm_train, so code that imports
the bare names fails at import instead of failing at run time with ValueError.
The top level keeps the det names (cake_rmsnorm_det, CakeRMSNormDetFunction,
cake_rmsnorm_det_train_forward / _backward / _backward_workspace_bytes), as in
the mHC and causal conv1d renames.

The names test asserts that the top level binds no non-det name and that
`from flashinfer import cake_rmsnorm` raises ImportError in a fresh interpreter;
the API test reaches the entry points through the module. docs/api/norm.rst
states the top-level rule.
@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_train.py tests/norm/test_cake_rmsnorm_train_all_sm.py tests/norm/test_cake_rmsnorm_train_api.py tests/norm/test_cake_rmsnorm_train_ckpt.py tests/norm/test_cake_rmsnorm_train_dist.py tests/norm/test_cake_rmsnorm_train_pairs.py tests/norm/test_cake_rmsnorm_train_stall.py tests/norm/test_cake_rmsnorm_train_vendored.py tests/norm/test_cake_rmsnorm_train_windows.py

This was referenced Nov 6, 2026
Merged
Open
averyhNV added 17 commits November 5, 2026 22:03
@averyhNV
docs(cake_rmsnorm): name the removed top-level names and their replac…
aadfa89 
…ements

docs/api/norm.rst lists the five removed top-level names one by one (the
cake_rmsnorm_train_* glob also matched the backend module, which stays) and
says what replaces them: the det names for the former deterministic behavior,
flashinfer.cake_rmsnorm_train.cake_rmsnorm with deterministic=False for the
non-deterministic package. The fresh-interpreter test matches the exact
ImportError text.
@averyhNV
Merge remote-tracking branch 'upstream/main' into aud-rmsnorm-names-v1
0ed952d
@averyhNV
refactor(flashinfer): one top-level lazy-export table for names serve…
d46f9b3 
…d from other modules

`flashinfer.__getattr__` served only the prims_ts names. Replace it with one
table, `_LAZY_EXPORTS: name -> (module, attribute)`, where an attribute of None
binds the module itself. The prims_ts names become rows of the table and behave
as before. Packages whose backend moves under `flashinfer.experimental` (issue
flashinfer-ai#8928) add their top-level names as rows, so `import flashinfer` does not import
the experimental tree. `tests/cpu/test_flashinfer_lazy_exports.py` checks that
every row is bound on first access only, that module rows bind modules, and that
unknown names still raise AttributeError.
@averyhNV
docs(experimental): Cake training backends under experimental/cake ba…
866fd69 
…ck stable facades and keep their sources in every wheel

flashinfer-ai#8928 moves each Cake training backend to flashinfer/experimental/cake/<module>/
behind a stable top-level facade. State in the experimental README that these
facades are not experimental APIs, load their backend lazily (facade functions
or module __getattr__, top-level names through flashinfer._LAZY_EXPORTS), and
that their csrc/** ships in every wheel. The release-wheel exclusion already
matches only flashinfer/experimental/<pkg>/csrc/**; say so next to the rule and
pin it in test_entry_rules for experimental/cake/<pkg>/csrc/**.
@averyhNV
refactor(cake_rmsnorm): mode-word names for the RMSNorm training pack…
b1414b7 
…ages (flashinfer-ai#8928)

The mode is part of the name, and only of the name.

- The non-deterministic facade flashinfer.cake_rmsnorm_train becomes
  flashinfer.cake_rmsnorm_nondet_train: cake_rmsnorm -> cake_rmsnorm_nondet,
  CakeRMSNormFunction -> CakeRMSNormNondetFunction,
  cake_rmsnorm_train_{forward,backward,backward_workspace,backward_workspace_bytes}
  -> cake_rmsnorm_nondet_train_*, is_cake_rmsnorm_train_supported ->
  is_cake_rmsnorm_nondet_train_supported, route_key ->
  cake_rmsnorm_nondet_train_route_key.
- The deterministic facade keeps its names; its helpers become
  cake_rmsnorm_det_train_{route_key,program_windows,weight_dtypes}.
- Both backends move under flashinfer.experimental.cake:
  flashinfer.experimental.cake.cake_rmsnorm_det_train and
  flashinfer.experimental.cake.cake_rmsnorm_nondet_train. The rendered files
  move unchanged except the build.json "package" field.
- No entry point takes `deterministic` any more (passing it is Python's
  TypeError); the ValueError path for values other than False is removed.
  Under torch.use_deterministic_algorithms(True) the non-deterministic
  backward still alerts (RuntimeError; warn_only: one warning, then its own
  backward) and names the det entry point; nothing else alerts.
- The old modules and names are removed without aliases, so old code fails
  at import. The flashinfer top level binds the five det entry points and
  the module cake_rmsnorm_nondet_train only.
- Tests, benchmarks and packaging (pyproject package-data, the manifest
  checker, the backend package table, labeler and pre-commit paths) follow;
  the names test checks the new contract.
@averyhNV
docs(norm): RMSNorm training mode rules and migration table (flashinf…
04ae7d3 
…er-ai#8928)

The two training RMSNorm sections name the new modules
(flashinfer.cake_rmsnorm_det_train, flashinfer.cake_rmsnorm_nondet_train) and
the renamed helpers. The "Mode and dispatch rules" note states that the mode
is part of the name only, that no entry point takes `deterministic`, the
torch.use_deterministic_algorithms rule and the top-level rule, and lists the
removed names with their replacements.
@averyhNV
style(cake_rmsnorm): two blank lines before the nondet backend's dete…
7b45cfc 
…rministic-mode check (flashinfer-ai#8928)

Removing the former flag helpers left one blank line between
SUPPORTED_COMPUTE_CAPABILITIES and the top-level
_check_deterministic_algorithms; ruff format wants two.
@averyhNV
Merge branch 'cake-lazy-exports-table' (flashinfer-ai#8932) into aud-…
9993467 
…rmsnorm-names-v1
@averyhNV
refactor(cake_rmsnorm): lazy facades and top-level rows for the exper…
8eb74f4 
…imental/cake backends (flashinfer-ai#8928)

The RMSNorm training backends live under flashinfer/experimental/cake/, and core
must not import that tree at module level (tests/experimental/test_containment.py
rules 2 and 3).

- flashinfer.cake_rmsnorm_det_train / flashinfer.cake_rmsnorm_nondet_train no
  longer import their backend at module level. The three entry points stay
  @flashinfer_api defs with unchanged signatures; their bodies call a
  functools.cache accessor that imports the backend on first use. The backend
  objects of __all__ (autograd class, constants, workspace / route / window /
  weight-dtype helpers, support query) resolve through the module __getattr__,
  are cached in the module and listed by __dir__; each is the backend's object.
- The entry points carry static docstrings: the backend function's written
  docstring, and for the forward / backward a Notes section naming the backend
  function whose docstring lists the program windows (the backend appends them
  from its registry at import). The facades no longer copy __doc__ from the
  backend at import. The facade module docstrings and the norm.rst
  non-deterministic section point at those backend notes.
- flashinfer/__init__.py: the five det re-exports and the nondet module binding
  become six rows of _LAZY_EXPORTS, so import flashinfer loads neither facade
  nor backend.
- Names test: facades may define __getattr__ (backends still may not); new
  checks for the six rows, lazy identity and caching, docstring parity with the
  backend, and fresh-interpreter module isolation (import flashinfer loads no
  RMSNorm module; importing a facade loads no backend; first use loads only its
  own backend).
@averyhNV
docs(cake_rmsnorm): no deterministic keyword in the nondet backward w…
ee8c224 
…indow notes (flashinfer-ai#8928)

The rendered program-selection notes of the non-deterministic package
(cake_windows.py NOTES, FORWARD_SECTION and BACKWARD_SECTION) still said
"A backward window (``deterministic=False``) runs one of two forms", naming a
keyword the package no longer takes. Drop the parenthetical at its three
occurrences; nothing else in the rendered file changes. The paired Kitchen
export template gets the identical edit, so a re-render reproduces this file.
@averyhNV
build: drop the csrc of listed experimental/cake backends from the re…
df49004 
…lease wheel

The release wheel (FLASHINFER_WHEEL_EXCLUDE_EXPERIMENTAL_SOURCES=1) drops
flashinfer/experimental/<package>/csrc/**. flashinfer-ai#8928 moves the Cake training
backends one level deeper, to flashinfer/experimental/cake/<module>/, where
the pattern does not match. Several of them were dropped at their old place;
moving them as is would add about 12 MB (compressed) to a wheel that is
93.7 MB today, over PyPI's 100 MB per-file cap.

_RELEASE_WHEEL_DROPPED_CAKE_BACKENDS lists the experimental/cake modules
whose csrc/** the flag also drops. It starts empty: each PR that moves a
backend that is dropped today adds its new module name as one row, so the
release wheel contents stay the same per package. Every other backend under
experimental/cake keeps its sources in every wheel.
@averyhNV
ci(api-diff): count _LAZY_EXPORTS rows as top-level exports
35edde1 
exported_names() read only the module-scope `from ... import` statements
of flashinfer/__init__.py. A name moved into a row of the lazy-export
table (flashinfer-ai#8932) was therefore reported as public_export_removed although
`flashinfer.<name>` still resolves. Rows now count like the import they
replace: ("<module>", "<attribute>") like `from <module> import
<attribute> as <name>`, ("<package>.<module>", None) like
`from <package> import <module> as <name>`. Rows spread in with `**` are
computed at runtime and are not read, as before.
@averyhNV
ci(api-diff): zip the lazy-export keys and rows strictly; ruff format…
bd95c0c 
… the new test
@averyhNV
test(lazy-exports): check that no row is bound before accessing any row
6b01fe8 
The test walked the rows in sorted order and accessed each one. An
attribute row from a module imports that module, and Python binds the
submodule on flashinfer, so a module-valued row for the same module that
sorts later was reported as "bound by import flashinfer" although
`import flashinfer` does not bind it. Check all rows first, then access.
@averyhNV
refactor(cake_rmsnorm): top-level row for the det facade module (flas…
57c3440 
…hinfer-ai#8928)

The eager det re-exports used to bind the submodule attribute
flashinfer.cake_rmsnorm_det_train as a side effect of import flashinfer.
With the RMSNorm names served by _LAZY_EXPORTS rows that attribute existed
only after something had imported the det facade.  Add the row
"cake_rmsnorm_det_train": (".cake_rmsnorm_det_train", None) next to the
nondet module row, so both facade modules are reachable from a bare
import flashinfer and the two modes behave the same.

The names test expects the seven rows, checks the det module row in a
fresh interpreter (it loads the det facade and no backend) and in the
flashinfer attribute import form.  docs/api/norm.rst lists both modules.
@averyhNV
docs(cake_rmsnorm): name the program-selection notes in the nondet ba…
44aa95d 
…ckward docstring (flashinfer-ai#8928)

"the form the notes below list for it" was true in the backend, whose
backward docstring gets the "Program selection by token count" notes
appended at import, but not in the facade's static copy, whose Notes
section only points to the backend function.  Name the notes instead, in
the backend's hand-written docstring (not rendered by the exporter) and in
the facade's copy, which the names test keeps byte-identical.
@averyhNV
Merge remote-tracking branch 'origin/main' into cake-release-wheel-dr…
70ebbfc 
…opped-backends

# Conflicts:
#	build_backend.py
#	flashinfer/experimental/README.md
#	tests/cpu/test_build_backend_slim_wheel.py
#	tests/cpu/test_flashinfer_lazy_exports.py
@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py tests/norm/test_cake_rmsnorm_nondet_train.py tests/norm/test_cake_rmsnorm_nondet_train_all_sm.py tests/norm/test_cake_rmsnorm_nondet_train_api.py tests/norm/test_cake_rmsnorm_nondet_train_ckpt.py tests/norm/test_cake_rmsnorm_nondet_train_dist.py tests/norm/test_cake_rmsnorm_nondet_train_pairs.py tests/norm/test_cake_rmsnorm_nondet_train_stall.py tests/norm/test_cake_rmsnorm_nondet_train_vendored.py tests/norm/test_cake_rmsnorm_nondet_train_windows.py tests/cpu/test_flashinfer_lazy_exports.py

@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py tests/norm/test_cake_rmsnorm_nondet_train.py tests/norm/test_cake_rmsnorm_nondet_train_all_sm.py tests/norm/test_cake_rmsnorm_nondet_train_api.py tests/norm/test_cake_rmsnorm_nondet_train_ckpt.py tests/norm/test_cake_rmsnorm_nondet_train_dist.py tests/norm/test_cake_rmsnorm_nondet_train_pairs.py tests/norm/test_cake_rmsnorm_nondet_train_stall.py tests/norm/test_cake_rmsnorm_nondet_train_vendored.py tests/norm/test_cake_rmsnorm_nondet_train_windows.py tests/cpu/test_flashinfer_lazy_exports.py

@yyihuang yyihuang changed the title refactor(cake_rmsnorm): det names only; the non-deterministic package requires deterministic=False refactor(cake_rmsnorm): mode-word names and experimental/cake backends (#8928) Nov 6, 2026
@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py tests/norm/test_cake_rmsnorm_nondet_train.py tests/norm/test_cake_rmsnorm_nondet_train_all_sm.py tests/norm/test_cake_rmsnorm_nondet_train_api.py tests/norm/test_cake_rmsnorm_nondet_train_ckpt.py tests/norm/test_cake_rmsnorm_nondet_train_dist.py tests/norm/test_cake_rmsnorm_nondet_train_pairs.py tests/norm/test_cake_rmsnorm_nondet_train_stall.py tests/norm/test_cake_rmsnorm_nondet_train_vendored.py tests/norm/test_cake_rmsnorm_nondet_train_windows.py tests/cpu/test_flashinfer_lazy_exports.py

@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py tests/norm/test_cake_rmsnorm_nondet_train.py tests/norm/test_cake_rmsnorm_nondet_train_all_sm.py tests/norm/test_cake_rmsnorm_nondet_train_api.py tests/norm/test_cake_rmsnorm_nondet_train_ckpt.py tests/norm/test_cake_rmsnorm_nondet_train_dist.py tests/norm/test_cake_rmsnorm_nondet_train_pairs.py tests/norm/test_cake_rmsnorm_nondet_train_stall.py tests/norm/test_cake_rmsnorm_nondet_train_vendored.py tests/norm/test_cake_rmsnorm_nondet_train_windows.py tests/cpu/test_flashinfer_lazy_exports.py

yyihuang added a commit that referenced this pull request Nov 6, 2026
@yyihuang @averyhNV
test(cake_package): check that every Cake package imports with flashi…
33491c8 
…nfer blocked (#8951)

Related issue: #8928
Requirement:
#8928 (comment)
(every backend is self-contained)

This PR adds one test file,
`tests/cpu/test_cake_packages_standalone.py`, and changes nothing else.
A build that vendors a Cake package (#6553) copies the package directory
into its own tree and imports it under its own name, without flashinfer.
This test is the only check that imports every Cake package with
flashinfer blocked, so the #8928 renames keep that layout working.

### What it checks

For each distinct import path in `PACKAGES` of
`tests/test_helpers/cake_package.py`, read when the test runs, a fresh
interpreter does the following:

1. Blocks flashinfer with `sys.modules["flashinfer"] = None` and a
meta-path finder that raises `ImportError` for `flashinfer` and
`flashinfer.*`.
2. Removes the working directory from `sys.path`, along with every entry
at or above a loaded directory. It keeps the other entries, so torch and
tvm_ffi still import.
3. Loads the package directory where it is, without a copy, as the
top-level package `vendored_<directory>`. A relative import that climbs
out of the directory therefore fails.
4. Imports every `.py` module of the directory one by one, including
modules in subdirectories. If `__init__.py` fails, its other modules are
still checked.
5. Fails on any import error. It also fails if a flashinfer module
leaked: any `flashinfer.*` entry left in `sys.modules`, or a
`flashinfer` entry other than the `None` sentinel.

A `PACKAGES` row fails if its directory or `__init__.py` does not exist.
The child interpreters run with `CUDA_VISIBLE_DEVICES` empty and `-B`,
and each has a 600 s timeout. The prebuilt/JIT loader `cake_build.py`
imports `flashinfer.jit` only inside a function (`_jit_toolchain`), so
it passes. The file also tests itself on temporary packages: 13 cases
for the child interpreter and 5 for sibling validation. That way a pass
still means something once the allowlist is empty.

Limits:
- Imports inside functions are not checked, because the test only
imports modules.
- An editable install also maps `b12x`, `build_backend` and
`build_utils` to the checkout. The test does not block these names, so a
package must not import them.

### Declared siblings

A KCP or dense nondet package that needs a single-rank package declares
it as a sibling:
- **Where siblings come from:** `build.json` `common.requires.packages`
(sibling directory names) when that key is present. Otherwise, the
`requires` field of the package's row in
`scripts/check_cake_build_manifests.py`.
- **Layout:** the package and its siblings load under one synthetic
parent, `vendored_<directory>.<directory>` beside
`vendored_<directory>.<sibling>`. This matches the layout of
`check_cake_build_manifests.py --vendored-gate`, so `from ..<sibling>
import x` resolves. Undeclared siblings cannot be imported.
- **Requirements:** each declared sibling must be a `PACKAGES` row with
the same parent package and an existing directory. Otherwise the package
fails.

### Burn-down `ALLOWLIST`

These four packages fail on main today (the `cake_dsa_indexer` alias row
left with its removal in #8936):

| import path | modules that fail today | reason |
|---|---|---|
| `flashinfer.experimental.cake_kda_det_train_kcp` | `cake_kcp`,
`cake_kcp_jit` | they import the single-rank package across directories
(`from ..cake_kda_det_train import`), undeclared |
| `flashinfer.experimental.cake_kda_train` | `__init__`, `cake_backend`,
`cake_build`, `cake_jit`, `cake_kcp`, `cake_kcp_jit`, `cake_kcp_plan` |
compatibility alias: `__init__.py` star-imports
`flashinfer.experimental.cake_kda_det_train`, and each module replaces
itself with a det package's module, imported by its flashinfer path |
| `flashinfer.experimental.cake_kda_train_kcp` | `cake_kcp`,
`cake_kcp_jit` | they import the single-rank package across directories
(`from ..cake.cake_kda_train import`), undeclared |
| `flashinfer.experimental.cake_stepfun_moe` | `__init__` |
`__init__.py` decorates its entry points with flashinfer's
`flashinfer_experimental_api` (`from ...api_logging import`) |

The allowlist is strict. An allowlisted package is marked xfail only
when exactly its listed modules fail. The test fails if a listed module
now imports or no longer exists, or if an unlisted module fails. A
leaked flashinfer module always fails the test, whether or not the
package is allowlisted. So the PR that fixes a package deletes its
modules from the entry in the same change, and deletes the entry once it
is empty. `test_allowlist_names_packages` fails if an entry's path is
not in `PACKAGES`, if an entry is duplicated or empty, or if `PACKAGES`
is empty.

How each row goes away:
- **The two KDA KCP rows:** the KDA rename PRs declare the single-rank
sibling (`common.requires.packages` in `build.json`, plus the README)
and resolve it from their parent package.
- **The alias row** (`cake_kda_train`): the alias is removed.
- **`cake_stepfun_moe`:** it imports flashinfer's `api_logging`. Its
owner decides how to vendor it.

### For the rename PRs

Each family updates only its own `PACKAGES` rows. The test reads the
table at run time, so no family copies this check. CI runs this file in
the unit-test run over `tests/`. The file lives in `tests/cpu` because
that run does not recurse into `tests/experimental`.

The `experimental-tests` block of a PR body accepts only
`tests/experimental/`. A PR that changes `PACKAGES`, a package directory
or a sibling declaration should therefore ask for a targeted run:
`@flashinfer-bot run tests/cpu/test_cake_packages_standalone.py`

### Tests

- Cluster CPU node, CUDA 13.4 development container, main 2e98955: 57
passed, 5 xfailed (exactly the five `ALLOWLIST` rows);
`scripts/check_cake_build_manifests.py`: 0 findings for 44 packages.
- Same check on the RMSNorm rename (#8910) merged into main: 57 passed,
5 xfailed; both renamed RMSNorm packages
(`flashinfer.experimental.cake.cake_rmsnorm_det_train`,
`flashinfer.experimental.cake.cake_rmsnorm_nondet_train`) import with
flashinfer blocked.
- Same check merged into main 555dd84 (after the conv1d, GDN and MoE
renames): 57 passed, 5 xfailed (same five `ALLOWLIST` rows).
- Current head merged into main 98fd797 (after the DSA indexer rename
#8936): 57 passed, 4 xfailed (exactly the four `ALLOWLIST` rows above);
`scripts/check_cake_build_manifests.py`: 0 findings for 44 packages.
- Same check merged into main 02b041c (after #8901, #8939 and #8922):
57 passed, 4 xfailed (same four `ALLOWLIST` rows);
`scripts/check_cake_build_manifests.py`: 0 findings for 44 packages.
- pre-commit: passes.

---------

Co-authored-by: Yingyi Huang <averyh@nvidia.com>
@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py tests/norm/test_cake_rmsnorm_nondet_train.py tests/norm/test_cake_rmsnorm_nondet_train_all_sm.py tests/norm/test_cake_rmsnorm_nondet_train_api.py tests/norm/test_cake_rmsnorm_nondet_train_ckpt.py tests/norm/test_cake_rmsnorm_nondet_train_dist.py tests/norm/test_cake_rmsnorm_nondet_train_pairs.py tests/norm/test_cake_rmsnorm_nondet_train_stall.py tests/norm/test_cake_rmsnorm_nondet_train_vendored.py tests/norm/test_cake_rmsnorm_nondet_train_windows.py tests/cpu/test_flashinfer_lazy_exports.py

@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py tests/norm/test_cake_rmsnorm_nondet_train.py tests/norm/test_cake_rmsnorm_nondet_train_all_sm.py tests/norm/test_cake_rmsnorm_nondet_train_api.py tests/norm/test_cake_rmsnorm_nondet_train_ckpt.py tests/norm/test_cake_rmsnorm_nondet_train_dist.py tests/norm/test_cake_rmsnorm_nondet_train_pairs.py tests/norm/test_cake_rmsnorm_nondet_train_stall.py tests/norm/test_cake_rmsnorm_nondet_train_vendored.py tests/norm/test_cake_rmsnorm_nondet_train_windows.py tests/cpu/test_flashinfer_lazy_exports.py

@yyihuang

yyihuang commented Nov 6, 2026

Copy link
Copy Markdown
Collaborator Author

@flashinfer-bot run tests/norm/test_cake_rmsnorm_det_train_names.py tests/norm/test_cake_rmsnorm_det_train.py tests/norm/test_cake_rmsnorm_det_train_api.py tests/norm/test_cake_rmsnorm_det_train_pairs.py tests/norm/test_cake_rmsnorm_det_train_stall.py tests/norm/test_cake_rmsnorm_det_train_training.py tests/norm/test_cake_rmsnorm_nondet_train.py tests/norm/test_cake_rmsnorm_nondet_train_all_sm.py tests/norm/test_cake_rmsnorm_nondet_train_api.py tests/norm/test_cake_rmsnorm_nondet_train_ckpt.py tests/norm/test_cake_rmsnorm_nondet_train_dist.py tests/norm/test_cake_rmsnorm_nondet_train_pairs.py tests/norm/test_cake_rmsnorm_nondet_train_stall.py tests/norm/test_cake_rmsnorm_nondet_train_vendored.py tests/norm/test_cake_rmsnorm_nondet_train_windows.py tests/cpu/test_flashinfer_lazy_exports.py

This branch has not been deployed

No deployments
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

At least 1 approving review is required to merge this pull request.

Assignees

No one assigned

Labels

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