From 1d5f7bbf37c7a1fb532984032d0ca12dc4ec9a8e Mon Sep 17 00:00:00 2001 From: Yingyi Huang Date: Thu, 5 Nov 2026 19:20:12 -0800 Subject: [PATCH] test(cake_gdn_train): stop asserting the deterministic packages' former names exist The non-deterministic GDN tests, bench and READMEs no longer depend on the former no-det names of the deterministic GDN packages (cake_gdn_det_train, cake_gdn_det_train_kcp) or on the former KDA package path (flashinfer.experimental.cake_kda_train), so they stay green before and after those names are removed: - test_cake_gdn_train_kcp_api.py: drop the test that the former names (flashinfer.gdn_train, flashinfer.experimental.cake_gdn_train and the module-level aliases of cake_gdn_det_train_kcp) are the det objects. Keep one absence test over a single list of the former spellings, now also on cake_backend, cake_build and cake_l2n, and a test that this package's entries and types are not the det package's (det names only). The KDA hygiene test refuses any import containing "cake_kda". - test_cake_gdn_train.py: compare this package's entries with the det package's det names instead of asserting its former aliases. - test_cake_gdn_train_kcp.py, test_cake_gdn_train_kcp_build.py: the KDA references use cake_kda_det_train_kcp.cake_kcp_plan (partition_det_kcp, build_context_det_kcp) and cake_kda_det_train.cake_build. - bench_cake_gdn_train_kcp.py: read LAST_CALL and graph_cache_clear of both routes through the package's cake_kcp module. - READMEs: the deterministic packages are named by their det names only. - cake_gdn_train_kcp provenance docstrings (cake_jit.py, cake_kcp.py, cake_kcp_plan.py) keep the historical KDA path at the cited commit and add the file's current det KDA path (text only). The former conv1d names used in the same files (bench, dist test, facade docstring) are migrated by the conv1d owner's own change. --- benchmarks/bench_cake_gdn_train_kcp.py | 15 +- .../cake/cake_gdn_train/README.md | 18 +- .../experimental/cake_gdn_train_kcp/README.md | 10 +- .../cake_gdn_train_kcp/cake_jit.py | 3 +- .../cake_gdn_train_kcp/cake_kcp.py | 3 +- .../cake_gdn_train_kcp/cake_kcp_plan.py | 5 +- tests/experimental/test_cake_gdn_train.py | 14 +- tests/experimental/test_cake_gdn_train_kcp.py | 22 +- .../test_cake_gdn_train_kcp_api.py | 211 +++++------------- .../test_cake_gdn_train_kcp_build.py | 4 +- 10 files changed, 104 insertions(+), 201 deletions(-) diff --git a/benchmarks/bench_cake_gdn_train_kcp.py b/benchmarks/bench_cake_gdn_train_kcp.py index 1176a541901..f70e51192e8 100644 --- a/benchmarks/bench_cake_gdn_train_kcp.py +++ b/benchmarks/bench_cake_gdn_train_kcp.py @@ -280,6 +280,8 @@ def __init__(self, kind: str) -> None: self._comm = self.api.gdn_comm_bytes_kcp self.handle_name = "RecomputeHandleKcp" self.workspace_name = "WorkspaceKcp" + # the host module of the route's package: its ``LAST_CALL`` records and ``graph_cache_clear`` live there + self.kcp = importlib.import_module(f"{self.pkg.__name__}.cake_kcp") def partition( self, cu: torch.Tensor, world: int, layout: str, cut_points: Any @@ -341,9 +343,7 @@ def workspace(self, ctx: Any, heads: tuple[int, int], device: torch.device) -> A def graph_cache_clear(self) -> None: """Release this route's captured graphs (before the next configuration and the process group's end).""" - fn = getattr(self.pkg, "graph_cache_clear", None) or getattr( - getattr(self.pkg, "cake_kcp", None), "graph_cache_clear", None - ) + fn = getattr(self.kcp, "graph_cache_clear", None) if fn is not None: fn() @@ -373,20 +373,19 @@ def bytes( } def modules(self) -> dict[str, Any]: - kcp = getattr(self.pkg, "cake_kcp", None) return { "facade": self.api.__name__, "facade_file": self.api.__file__, "package": self.pkg.__name__, "package_file": self.pkg.__file__, - "cake_kcp_file": getattr(kcp, "__file__", None), + "cake_kcp_file": self.kcp.__file__, } def clear_last_call(self) -> None: - self.pkg.LAST_CALL.clear() + self.kcp.LAST_CALL.clear() def last_call(self, kind: str) -> Optional[dict[str, Any]]: - rec = self.pkg.LAST_CALL.get(kind) + rec = self.kcp.LAST_CALL.get(kind) if not rec: return None out = { @@ -786,7 +785,7 @@ def _probe( torch.cuda.synchronize(device) mem["after_bytes"] = torch.cuda.memory_allocated(device) - base mem["step_peak_bytes"] = max(mem["fwd_peak_bytes"], mem.get("bwd_peak_bytes", 0)) - ran = sorted(kind for kind, route in routes.items() if route.pkg.LAST_CALL) + ran = sorted(kind for kind, route in routes.items() if route.kcp.LAST_CALL) if sota: made = {key: lib.calls[key] - before[key] for key in lib.calls} need = (SOTA_GDN_CALL[arm],) + (("conv",) if len(grads_out) == 2 else ()) diff --git a/flashinfer/experimental/cake/cake_gdn_train/README.md b/flashinfer/experimental/cake/cake_gdn_train/README.md index 870f8cb90f8..f18e0699375 100644 --- a/flashinfer/experimental/cake/cake_gdn_train/README.md +++ b/flashinfer/experimental/cake/cake_gdn_train/README.md @@ -85,15 +85,19 @@ non-deterministic one). No launch of this package is unordered, so PyTorch's deterministic algorithms (`torch.use_deterministic_algorithms(True)`, with or without `warn_only`) -change nothing: no entry raises PyTorch's alert. The deterministic packages, -their names, defaults and aliases are unchanged. +change nothing: no entry raises PyTorch's alert. -| name | what it is now | +The deterministic packages are named by their det names only: + +| name | what it is | |---|---| -| `flashinfer.experimental.cake_gdn_det_train` and its aliases `chunk_gated_delta_rule`, `ChunkGatedDeltaRuleFunction`, `gdn_chunk_forward`, `gdn_chunk_backward` | unchanged: the deterministic one-rank package | -| `flashinfer.experimental.cake_gdn_train` (the KCP package of FlashInfer v0.7.3rc1) | unchanged: a legacy alias of the deterministic KCP package `flashinfer.experimental.cake_gdn_det_train_kcp` | -| `flashinfer.gdn_train`, `flashinfer.cake_gdn_det_train_kcp` | unchanged: deterministic KCP entry points | -| `flashinfer.experimental.cake.cake_gdn_train` | new: this non-deterministic one-rank package | +| `flashinfer.experimental.cake_gdn_det_train` (`chunk_gated_delta_rule_det`, `ChunkGatedDeltaRuleDetFunction`, `gdn_chunk_det_forward`, `gdn_chunk_det_backward`) | the deterministic one-rank package | +| `flashinfer.experimental.cake_gdn_det_train_kcp`, `flashinfer.cake_gdn_det_train_kcp` | the deterministic KCP package and its entry points | +| `flashinfer.experimental.cake.cake_gdn_train` | this non-deterministic one-rank package | + +The no-det spellings `chunk_gated_delta_rule`, `ChunkGatedDeltaRuleFunction`, +`gdn_chunk_forward` and `gdn_chunk_backward` name this package's +non-deterministic entries; a deterministic call uses the det names above. ## Chain diff --git a/flashinfer/experimental/cake_gdn_train_kcp/README.md b/flashinfer/experimental/cake_gdn_train_kcp/README.md index f8c10d72aae..11328d55ba5 100644 --- a/flashinfer/experimental/cake_gdn_train_kcp/README.md +++ b/flashinfer/experimental/cake_gdn_train_kcp/README.md @@ -76,12 +76,10 @@ Every name of this package is the deterministic twin's name without `det`: | `ChunkGDNDetFunctionKcp` | `ChunkGDNFunctionKcp` | | `packed_qkv_views_det_kcp`, `PackedQKVSplitDetKcp` | `packed_qkv_views_kcp`, `PackedQKVSplitKcp` (package level) | -The legacy names of the deterministic package keep meaning the deterministic -objects: `flashinfer.gdn_train.chunk_gdn_kcp`, `KcpContext`, ... and the -legacy package path `flashinfer.experimental.cake_gdn_train` (with its -`chunk_gdn_kcp`) are aliases of `cake_gdn_det_train_kcp` objects, never of -this package's. The name `chunk_gdn_kcp` of this package is a different -object from the legacy alias of the same spelling. +The deterministic package is named by its det names (left column) only. +This package binds none of the deterministic package's former no-det names +(such as `kcp_partition` or `KcpContext`), and `chunk_gdn_kcp` is this +package's own non-deterministic entry, never a deterministic object. The generated libraries are named `cake_gdn_train_<20 hex>` for the per-chunk programs and `cake_gdn_train_<20 hex>_kcp` for the diff --git a/flashinfer/experimental/cake_gdn_train_kcp/cake_jit.py b/flashinfer/experimental/cake_gdn_train_kcp/cake_jit.py index f8dd41f0a3b..729a13eb614 100644 --- a/flashinfer/experimental/cake_gdn_train_kcp/cake_jit.py +++ b/flashinfer/experimental/cake_gdn_train_kcp/cake_jit.py @@ -91,7 +91,8 @@ ``K K^T`` Gram, no block inverse): every other output bit for bit ``fwd_prep_l2n``'s, ``akk`` read, not written; grid (NC, H_qk), two resident CTAs per SM as #1 -Provenance of the rules: ``flashinfer/experimental/cake_kda_train/cake_kcp_jit.py`` (records_of_program, +Provenance of the rules: ``flashinfer/experimental/cake_kda_train/cake_kcp_jit.py`` (now +``flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_jit.py``; records_of_program, select_record, launch_grid; sha256 c405e7dcd36c9cbe0ff01e3e3194fd760491504655bda28a646c2067a23b4c90) at FlashInfer 13d7fcf4041635f5b7f573bb0dbab23570ae1034, adapted to the GDN vocabulary; the loader glue moved to :mod:`.cake_backend`. diff --git a/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py b/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py index 4cca959e06d..098fdfec002 100644 --- a/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py +++ b/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp.py @@ -19,7 +19,8 @@ :class:`SavedKcp`, the autograd Function and the entry points ``chunk_gdn_kcp`` / ``chunk_gdn_fwd_kcp`` / ``chunk_gdn_bwd_kcp``. -Provenance: adapted from ``flashinfer/experimental/cake_kda_train/cake_kcp.py`` +Provenance: adapted from ``flashinfer/experimental/cake_kda_train/cake_kcp.py`` (now +``flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp.py``) at FlashInfer ``13d7fcf4041635f5b7f573bb0dbab23570ae1034`` (the KDA KCP host module). Verbatim parts: the compact pools and chains (``Indexed``, ``Chains``), the exchange tables (``_Tables``), the per-call state diff --git a/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp_plan.py b/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp_plan.py index 5c4f3f79272..1400fa0553e 100644 --- a/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp_plan.py +++ b/flashinfer/experimental/cake_gdn_train_kcp/cake_kcp_plan.py @@ -13,11 +13,12 @@ See the License for the specific language governing permissions and limitations under the License. -Provenance: package-local copy of ``flashinfer/experimental/cake_kda_train/cake_kcp_plan.py`` at FlashInfer +Provenance: package-local copy of ``flashinfer/experimental/cake_kda_train/cake_kcp_plan.py`` (now +``flashinfer/experimental/cake_kda_det_train_kcp/cake_kcp_plan.py``) at FlashInfer commit 13d7fcf4041635f5b7f573bb0dbab23570ae1034 (sha256 0f2d8b06bbea49df37b043fba3026a06799f0a0cba09d697e490d02bb747dfbf), adapted for the GDN KCP route: the head count of every per-head state is the value-head count ``HV`` (``H := HV``), the plan rule is a keyword-only argument (no environment variable), and the byte queries live -in ``cake_bytes`` (GDN saved set and exchange formulas). ``cake_gdn_train_kcp`` never imports ``cake_kda_train``; +in ``cake_bytes`` (GDN saved set and exchange formulas). ``cake_gdn_train_kcp`` never imports a KDA package; a CPU test asserts field-identical partitions and contexts on the same inputs. Metadata layer of the context-parallel (KCP) route of the Cake packed GDN diff --git a/tests/experimental/test_cake_gdn_train.py b/tests/experimental/test_cake_gdn_train.py index d75fc914d31..e0cba11eb9f 100644 --- a/tests/experimental/test_cake_gdn_train.py +++ b/tests/experimental/test_cake_gdn_train.py @@ -1820,15 +1820,11 @@ def test_the_package_exports_the_det_free_names_only() -> None: for name in DET_NAMES: for module in (cake_gdn_train, cake_chain): assert not hasattr(module, name) and name not in module.__all__, name - # the deterministic package keeps its names and its aliases, distinct objects from this package's entries - assert ( - cake_gdn_det_train.chunk_gated_delta_rule - is cake_gdn_det_train.chunk_gated_delta_rule_det - ) - assert ( - cake_gdn_train.chunk_gated_delta_rule - is not cake_gdn_det_train.chunk_gated_delta_rule - ) + # the deterministic package's entries (its det names) are distinct objects from this package's entries + for name, det_name in zip(ENTRY_NAMES, DET_NAMES, strict=True): + assert getattr(cake_gdn_train, name) is not getattr( + cake_gdn_det_train, det_name + ), name assert cake_jit.PACKAGE == "flashinfer.experimental.cake.cake_gdn_train" assert cake_jit.LIBRARY_PREFIX == "cake_gdn_train" diff --git a/tests/experimental/test_cake_gdn_train_kcp.py b/tests/experimental/test_cake_gdn_train_kcp.py index 22f1990ea97..111499780a0 100644 --- a/tests/experimental/test_cake_gdn_train_kcp.py +++ b/tests/experimental/test_cake_gdn_train_kcp.py @@ -15,8 +15,8 @@ CPU tests of the GDN KCP metadata layer of the NON-DETERMINISTIC package ``flashinfer.experimental.cake_gdn_train_kcp`` (the twin of ``cake_gdn_det_train_kcp`` with the same contract and accuracy budgets, not its bitwise reproducibility): -the package-local plan module against ``cake_kda_train.cake_kcp_plan`` (field-identical partitions and contexts on the -step-0 rows with ``H := HV``; the KDA package is imported HERE only, never by the package), the per-rank crossing +the package-local plan module against ``cake_kda_det_train_kcp.cake_kcp_plan`` (field-identical partitions and contexts on +the step-0 rows with ``H := HV``; the KDA package is imported HERE only, never by the package), the per-rank crossing counts and exchange bytes against the rows of record, the table builders (pool addressing, segment lists, chains), head ownership incl. the padded ``(HV, W) = (3, 2)`` case, the byte queries, the registry seed, an emulated multi-rank row (W ranks on threads), and the package's host orchestration on the fake backend @@ -381,14 +381,14 @@ def _sha(obj) -> str: @ROWS @pytest.mark.parametrize("path", ROW_FILES, ids=ROW_IDS) def test_plan_parity_with_kda_on_rows(path, monkeypatch) -> None: - """Result 1: ``cake_gdn_train_kcp.cake_kcp_plan`` is field-identical to ``cake_kda_train.cake_kcp_plan`` on every - step-0 row (partition, every rank's context and head plan at ``H := HV``) and reproduces the row's recorded + """Result 1: ``cake_gdn_train_kcp.cake_kcp_plan`` is field-identical to ``cake_kda_det_train_kcp.cake_kcp_plan`` on + every step-0 row (partition, every rank's context and head plan at ``H := HV``) and reproduces the row's recorded fragment table. Result 2: ``n_in`` / ``n_out`` / ``max_chain_fragments`` and the exchange bytes of ``gdn_comm_bytes_kcp`` equal the row JSON. Both modules plan under the per-rank rule they share, ``"wave_fill1"`` (the GDN default is the row rule, the KDA default its own tie rule; each is selected explicitly here: the keyword-only ``p_rule`` on the GDN side, the private test hook on the KDA side).""" kp = importlib.import_module( - "flashinfer.experimental.cake_kda_train.cake_kcp_plan" + "flashinfer.experimental.cake_kda_det_train_kcp.cake_kcp_plan" ) # the test only # the KDA module accepts its private test hook only after its opt-in (enable_test_knobs) monkeypatch.setattr(kp, "_TEST_KNOBS_ENABLED", True) @@ -398,7 +398,7 @@ def test_plan_parity_with_kda_on_rows(path, monkeypatch) -> None: W = int(row["cp"]) args = _partition_inputs(row) part = gp.partition_kcp(**args) - kpart = kp.kcp_partition(**args) + kpart = kp.partition_det_kcp(**args) assert part.to_dict() == kpart.to_dict(), "partition differs from the KDA module" assert _plain(part.to_dict()["fragments"]) == row["fragments"], ( "partition differs from the row of record" @@ -419,7 +419,7 @@ def test_plan_parity_with_kda_on_rows(path, monkeypatch) -> None: ctx = gp.build_context_kcp( part, None, r, CPU, 4, sm_count=SM_COUNT, p_rule="wave_fill1" ) # the per-rank rule both modules share; the GDN default "row_cost" is tested on the rows below - kctx = kp.build_kcp_context( + kctx = kp.build_context_det_kcp( kpart, None, r, @@ -2037,8 +2037,8 @@ def test_registry_seed() -> None: def test_package_imports_no_kda_and_no_environ() -> None: - """The package never imports ``cake_kda_train`` and reads no environment variable; the plan - module imports torch, torch.distributed and dataclasses only.""" + """The package never imports a KDA package (no module path or name containing ``cake_kda``) and reads no + environment variable; the plan module imports torch, torch.distributed and dataclasses only.""" for path in sorted(PACKAGE_DIR.glob("*.py")): tree = ast.parse(path.read_text()) for node in ast.walk(tree): @@ -2048,8 +2048,8 @@ def test_package_imports_no_kda_and_no_environ() -> None: if isinstance(node, ast.ImportFrom) and node.module else [] ) - assert not any("cake_kda_train" in n for n in names), ( - f"{path.name} imports cake_kda_train" + assert not any("cake_kda" in n for n in names), ( + f"{path.name} imports a KDA package" ) assert "os" not in [ a.name for a in node.names if isinstance(node, ast.Import) diff --git a/tests/experimental/test_cake_gdn_train_kcp_api.py b/tests/experimental/test_cake_gdn_train_kcp_api.py index 1e43736fb46..764452a4466 100644 --- a/tests/experimental/test_cake_gdn_train_kcp_api.py +++ b/tests/experimental/test_cake_gdn_train_kcp_api.py @@ -16,10 +16,10 @@ signatures, the facade's contract names, the numerical mode (``deterministic`` on every entry: the package refuses everything but ``False`` and, every path being fixed-order, raises no alert under ``torch.use_deterministic_algorithms``; the facade routes ``True`` and torch's deterministic mode to the twin, warns once, refuses a non-bool and a saved set of -the other route; no context, partition or options field holds the mode), the names (the twin's legacy aliases stay -the twin's objects; this package copies none of them and no public name carries a mode), the package hygiene rules -(no environment variable, no ``cake_kda_train`` import, standalone: no import of another FlashInfer package, its twin -included, and no library name of the twin) and the host tables built on the CPU for the correctness rows' +the other route; no context, partition or options field holds the mode), the names (this package and the facade copy +none of the twin's pre-rename names and share no entry or type with it; no public name carries a mode), the package +hygiene rules (no environment variable, no KDA package import, standalone: no import of another FlashInfer package, +its twin included, and no library name of the twin) and the host tables built on the CPU for the correctness rows' partitions. No kernel runs here: the compute entries run on the fake backend (the CPU emulation), directly or through ``KernelBackend`` with a stub launch glue (the architecture the kernel backend passes to the #7 walks). """ @@ -663,141 +663,72 @@ def test_package_init_re_exports_the_route() -> None: assert n in pkg.__all__ -# The deterministic twin's pre-rename names: legacy facade name -> det canonical facade name, legacy package-only -# name -> det canonical name, and the module-level aliases at the ends of the det package's modules. None of them is -# copied into this package; ``chunk_gdn_kcp`` is the one spelling both carry (a det legacy alias there, this package's -# canonical entry here). -LEGACY_FACADE_NAMES = { - "kcp_partition": "partition_det_kcp", - "build_kcp_context": "build_context_det_kcp", - "chunk_gdn_kcp": "chunk_gdn_det_kcp", - "chunk_gdn_kcp_fwd": "chunk_gdn_det_fwd_kcp", - "chunk_gdn_kcp_bwd": "chunk_gdn_det_bwd_kcp", - "gdn_kcp_workspace_bytes": "gdn_det_workspace_bytes_kcp", - "gdn_kcp_saved_bytes": "gdn_det_saved_bytes_kcp", - "gdn_kcp_comm_bytes": "gdn_det_comm_bytes_kcp", - "is_gdn_train_supported": "is_gdn_det_train_kcp_supported", - "KcpOptions": "OptionsDetKcp", - "KcpContext": "ContextDetKcp", - "KcpPartition": "PartitionDetKcp", - "KcpRecomputeHandle": "RecomputeHandleDetKcp", - "KcpSaved": "SavedDetKcp", - "KcpWorkspace": "WorkspaceDetKcp", -} -LEGACY_PACKAGE_NAMES = { - **{ - k: v - for k, v in LEGACY_FACADE_NAMES.items() - if k not in ("is_gdn_train_supported",) - }, - "ChunkGDNKcpFunction": "ChunkGDNDetFunctionKcp", - "KcpCall": "CallDetKcp", - "KcpFragment": "FragmentDetKcp", -} -LEGACY_MODULE_ALIASES = { - "cake_kcp": { - "ChunkGDNKcpFunction": "ChunkGDNDetFunctionKcp", - "KcpCall": "CallDetKcp", - "KcpOptions": "OptionsDetKcp", - "KcpRecomputeHandle": "RecomputeHandleDetKcp", - "KcpSaved": "SavedDetKcp", - "KcpWorkspace": "WorkspaceDetKcp", - "chunk_gdn_kcp": "chunk_gdn_det_kcp", - "chunk_gdn_kcp_bwd": "chunk_gdn_det_bwd_kcp", - "chunk_gdn_kcp_fwd": "chunk_gdn_det_fwd_kcp", - "gdn_kcp_comm_bytes": "gdn_det_comm_bytes_kcp", - "gdn_kcp_saved_bytes": "gdn_det_saved_bytes_kcp", - "gdn_kcp_workspace_bytes": "gdn_det_workspace_bytes_kcp", - }, - "cake_kcp_plan": { - "KcpContext": "ContextDetKcp", - "KcpFragment": "FragmentDetKcp", - "KcpPartition": "PartitionDetKcp", - "build_kcp_context": "build_context_det_kcp", - "kcp_partition": "partition_det_kcp", - }, - "cake_bytes": { - "gdn_kcp_comm_bytes": "gdn_det_comm_bytes_kcp", - "gdn_kcp_saved_bytes": "gdn_det_saved_bytes_kcp", - "gdn_kcp_workspace_bytes": "gdn_det_workspace_bytes_kcp", - "gdn_kcp_saved_layout": "gdn_det_saved_layout_kcp", - "gdn_kcp_workspace_layout": "gdn_det_workspace_layout_kcp", - }, - "cake_options": {"KcpOptions": "OptionsDetKcp"}, -} -SHARED_SPELLINGS = { - "chunk_gdn_kcp" -} # a det legacy alias and this package's canonical entry +# The deterministic twin's pre-rename spellings (its former facade, package and module-level names). This package +# and the facade copy none of them. ``chunk_gdn_kcp`` is not listed: it is this package's canonical entry. +LEGACY_DET_SPELLINGS = frozenset( + { + "kcp_partition", + "build_kcp_context", + "chunk_gdn_kcp_fwd", + "chunk_gdn_kcp_bwd", + "gdn_kcp_workspace_bytes", + "gdn_kcp_saved_bytes", + "gdn_kcp_comm_bytes", + "gdn_kcp_saved_layout", + "gdn_kcp_workspace_layout", + "is_gdn_train_supported", + "ChunkGDNKcpFunction", + "KcpCall", + "KcpContext", + "KcpFragment", + "KcpOptions", + "KcpPartition", + "KcpRecomputeHandle", + "KcpSaved", + "KcpWorkspace", + } +) SUBMODULES = ( + "cake_backend", + "cake_build", "cake_bytes", "cake_fake_backend", "cake_jit", "cake_kcp", "cake_kcp_plan", + "cake_l2n", "cake_options", "cake_tables", ) -def _nondet_objects() -> list: - """Every public object of this package (its ``__all__``, every function / class its modules define) and of the - facade (its ``__all__`` and its entries).""" - pkg, facade = _mod(), importlib.import_module(FACADE) - objs = [getattr(pkg, n) for n in pkg.__all__] - objs += [getattr(facade, n) for n in facade.__all__] - for sub in SUBMODULES: - mod = _mod(sub) - objs += [ - v - for v in vars(mod).values() - if callable(v) and str(getattr(v, "__module__", "")).startswith(PACKAGE) - ] - return objs +def test_no_module_of_the_package_or_the_facade_binds_a_legacy_det_spelling() -> None: + """No pre-rename spelling of the twin was copied: the package, each of its modules and the facade neither bind + nor export any of them. ``chunk_gdn_kcp`` is this package's own entry, defined in ``cake_kcp`` under that name.""" + facade = importlib.import_module(FACADE) + pkg = _mod() + for module in [pkg, facade] + [_mod(sub) for sub in SUBMODULES]: + exported = set(getattr(module, "__all__", ())) + for name in sorted(LEGACY_DET_SPELLINGS): + assert not hasattr(module, name), (module.__name__, name) + assert name not in exported, (module.__name__, name) + entry = _mod("cake_kcp").chunk_gdn_kcp + assert pkg.chunk_gdn_kcp is entry + assert ( + entry.__module__ == f"{PACKAGE}.cake_kcp" and entry.__name__ == "chunk_gdn_kcp" + ) -def test_legacy_names_are_the_deterministic_objects_never_the_non_deterministic_ones() -> ( - None -): - """``flashinfer.gdn_train``, ``flashinfer.experimental.cake_gdn_train`` and the det package's module-level aliases - stay the deterministic twin's compatibility paths: every legacy name IS its det object and is NOT any object of - this package or of the ``flashinfer.cake_gdn_train_kcp`` facade (in particular ``chunk_gdn_kcp``, spelled by both).""" - det_facade = importlib.import_module(DET_FACADE) +def test_the_package_shares_no_entry_or_type_with_the_twin() -> None: + """This package's entries (also through the facade) and types are never the twin's objects (the twin's det names).""" det_pkg = importlib.import_module(DET_PACKAGE) - nondet = _nondet_objects() - assert len(nondet) > 50 - - def is_nondet(obj) -> bool: - return any(obj is o for o in nondet) - - legacy = importlib.import_module("flashinfer.gdn_train") - for old, new in LEGACY_FACADE_NAMES.items(): - obj = getattr(legacy, old) - assert obj is getattr(det_facade, new), old - assert not is_nondet(obj), old - legacy_pkg = importlib.import_module("flashinfer.experimental.cake_gdn_train") - for old, new in LEGACY_PACKAGE_NAMES.items(): - obj = getattr(legacy_pkg, old) - assert obj is getattr(det_pkg, new), old - assert not is_nondet(obj), old - for sub, aliases in LEGACY_MODULE_ALIASES.items(): - det_mod = importlib.import_module(f"{DET_PACKAGE}.{sub}") - for old, new in aliases.items(): - obj = getattr(det_mod, old) - assert obj is getattr(det_mod, new), (sub, old) - assert not is_nondet(obj), (sub, old) - # the shared spelling: the det alias, never this package's entry or the facade's dispatcher facade = importlib.import_module(FACADE) - assert legacy.chunk_gdn_kcp is det_facade.chunk_gdn_det_kcp - assert legacy_pkg.chunk_gdn_kcp is det_pkg.chunk_gdn_det_kcp for nondet_entry in ( _mod().chunk_gdn_kcp, _mod("cake_kcp").chunk_gdn_kcp, facade.chunk_gdn_kcp, ): - assert nondet_entry is not legacy.chunk_gdn_kcp - assert nondet_entry is not legacy_pkg.chunk_gdn_kcp assert nondet_entry is not det_pkg.chunk_gdn_det_kcp - # and this package's types are not the twin's for name in ( "OptionsKcp", "ContextKcp", @@ -807,32 +738,7 @@ def is_nondet(obj) -> bool: ): assert getattr(_mod(), name) is not getattr( det_pkg, name.replace("Kcp", "DetKcp") - ) - - -def test_the_package_exports_no_legacy_alias_name() -> None: - """No pre-rename name of the twin was copied: the package, its modules and the facade bind none of them (except - ``chunk_gdn_kcp``, this package's canonical entry, defined in ``cake_kcp`` under its own name).""" - legacy_names = ( - set(LEGACY_FACADE_NAMES) - | set(LEGACY_PACKAGE_NAMES) - | {n for aliases in LEGACY_MODULE_ALIASES.values() for n in aliases} - ) - facade = importlib.import_module(FACADE) - pkg = _mod() - for name in sorted(legacy_names - SHARED_SPELLINGS): - assert not hasattr(pkg, name), name - assert name not in pkg.__all__, name - assert not hasattr(facade, name), name - for sub in SUBMODULES: - assert not hasattr(_mod(sub), name), (sub, name) - entry = _mod("cake_kcp").chunk_gdn_kcp - assert pkg.chunk_gdn_kcp is entry - assert ( - entry.__module__ == f"{PACKAGE}.cake_kcp" and entry.__name__ == "chunk_gdn_kcp" - ) - for module in [pkg] + [_mod(sub) for sub in SUBMODULES]: - assert not hasattr(module, "KcpContext"), module.__name__ + ), name def test_no_public_name_of_the_package_or_the_facade_names_a_mode() -> None: @@ -1732,7 +1638,7 @@ def test_the_registry_names_no_library_of_the_deterministic_twin() -> None: # ---------------------------------------------------------------------------------------------------------------- -# package hygiene (no cake_kda_train import; no os.environ; the host module's import allow-list) +# package hygiene (no KDA package import; no os.environ; the host module's import allow-list) def _package_sources() -> list[pathlib.Path]: @@ -1746,25 +1652,22 @@ def _facade_source() -> pathlib.Path: return pathlib.Path(importlib.import_module(FACADE).__file__).resolve() -def test_package_never_imports_cake_kda_train() -> None: +def test_package_never_imports_a_kda_package() -> None: + """Neither the package nor the facade imports a KDA package: no imported module path or name contains ``cake_kda`` + (the deterministic and the non-deterministic KDA packages).""" for path in _package_sources() + [_facade_source()]: tree = ast.parse(path.read_text()) for node in ast.walk(tree): if isinstance(node, ast.Import): for a in node.names: - assert "cake_kda_train" not in a.name, (path.name, a.name) + assert "cake_kda" not in a.name, (path.name, a.name) elif isinstance(node, ast.ImportFrom): - assert "cake_kda_train" not in (node.module or ""), ( + assert "cake_kda" not in (node.module or ""), ( path.name, node.module, ) for a in node.names: - assert a.name != "cake_kda_train", (path.name, a.name) - assert ( - "cake_kda_train" - not in path.read_text().replace("Provenance", "").split("import os")[0] - or True - ) + assert "cake_kda" not in a.name, (path.name, a.name) def test_package_reads_no_environment_variable() -> None: diff --git a/tests/experimental/test_cake_gdn_train_kcp_build.py b/tests/experimental/test_cake_gdn_train_kcp_build.py index 771c11c46bd..b8ea827abab 100644 --- a/tests/experimental/test_cake_gdn_train_kcp_build.py +++ b/tests/experimental/test_cake_gdn_train_kcp_build.py @@ -278,7 +278,7 @@ def test_the_package_imports_nothing_from_the_kda_training_package() -> None: names = [node.module or ""] + [alias.name for alias in node.names] else: continue - if any("cake_kda_train" in name for name in names): + if any("cake_kda" in name for name in names): offenders.append(f"{path.name}:{node.lineno}") assert not offenders, offenders @@ -328,7 +328,7 @@ def test_the_loader_and_header_are_the_pinned_bytes_of_every_cake_package() -> N pytest.skip("csrc/cake_ffi_checks.h is not rendered into the package yet") kda_dir = Path( importlib.import_module( - "flashinfer.experimental.cake_kda_train.cake_build" + "flashinfer.experimental.cake_kda_det_train.cake_build" ).__file__ ).parent for relative in ("cake_build.py", "csrc/cake_ffi_checks.h"):