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"):