chore(ci): kernel dtype-key checker cannot see launches that move into headers
- Dominant language
- Rust
- Stars
- 467
- Forks
- 54
- Avg merge
- 4h 25m
- Merged PRs (30d)
- 310
Description
Part of #1801
## Problem / Background
`scripts/ci/check_kernel_dtype_keys.py` decides what to check by two rules that are both invisible from the files being checked, so an ordinary refactor can silently empty its scope while the check still reports OK.
The rule it enforces is the #1053 and #1054 class: on CUDA, `custom_kernel.cpp` builds the JIT cache name from `name + template_arguments_hash(template_args)` and nothing else, so a launch whose `template_args` are all ints hashes to one name for every input dtype, and whichever dtype compiles first wins for the process. Metal's key already carries the dtypes. This check is the only place the omission is caught without CUDA hardware (`.github/workflows/ci.yml:285-296`).
## Current Behavior
Two invisible scope rules:
1. `main()` at `scripts/ci/check_kernel_dtype_keys.py:98-105` globs `*.cpp` only, over `SEARCH_DIRS = ("src/lib/mlx-cpp/turbo", "src/lib/mlxcel-core/cpp")`. A JIT launch that moves into a header is never scanned, with or without the token.
2. `check_file` at `scripts/ci/check_kernel_dtype_keys.py:78-80` returns early on `if "cuda_kernel(" not in src`. A launch that moves out of these files into a shared helper takes the whole file out of scope.
Eight files are in scope today, and they are the same eight that #1803 touched (via PR #1869): `src/lib/mlx-cpp/turbo/fused_norm.cpp`, `fused_rope_append.cpp`, `paged_attention.cpp`, `paged_attention_v2.cpp`, `paged_attention_v2_merge.cpp`, `sampling.cpp`, `sampling_rejection.cpp`, and `src/lib/mlxcel-core/cpp/mlx_cxx_kernels.cpp`. That change moved only the backend-selection boolean into `turbo/gpu_backend.cpp` and left every launch where it was, so the scope held, but it held by luck rather than by design: the same refactor carried one step further would have dropped all eight files and the check would still have printed OK.
The reported count is not the in-scope count. `main()` increments `scanned` for every `*.cpp` under the search directories before `check_file` applies the token filter, so the output currently says 16 while 8 files are actually checked, and adding an unrelated `.cpp` raises the number without changing what is checked.
`docs/code-guidelines.md:164` documents the token-based scoping as deliberate ("a Metal-only launcher is out of scope until someone adds a CUDA port to it, at which point the check starts applying on its own"). The intent is sound; the failure mode is that scope can also shrink to zero without anyone noticing.
## Scope
**In scope:** `scripts/ci/check_kernel_dtype_keys.py`, a new companion test, and the scoping paragraph in `docs/code-guidelines.md:164`.
**Out of scope:** changing the dtype-keying rule itself, or adding an allowlist (the script's docstring explains why there deliberately is none). Rewriting `TEMPLATE_ARGS_RE` or `ENTRY_RE`.
## Proposed Solution
Make the scope robust to a refactor, and make a scope that has shrunk visible rather than silent. Options to weigh:
1. Extend the glob to `*.h`, `*.hpp`, and `*.cuh` alongside `*.cpp`, so a launch that moves into a header stays in scope.
2. Track the in-scope file list separately from the scanned count and print both, so the success line reads as in-scope over scanned rather than a single number that only ever grows.
3. Pin the expected in-scope set (a constant in the script or a small checked-in manifest) and fail when a file leaves it without an accompanying update, the same way a snapshot test fails on an unexplained drop.
Options 1 and 2 are cheap and complementary; option 3 is what actually catches the "moved into a shared helper" case, at the cost of one more thing to update when a launcher is legitimately retired.
## Implementation Notes
- **Reuse**: the repository already has two companion tests for CI scripts in exactly this shape, `scripts/ci/check_cross_repo_refs_test.sh` and `scripts/ci/check_llama_compat_manifest_test.sh`, both building a temporary tree and asserting the script's exit code. Follow one of them rather than inventing a new harness. The Python-side alternative is `tests/test_quality_gate_script.py`.
- **Wiring**: `make verify-kernel-dtype-keys` (`Makefile:735-738`) and the `kernel dtype keys` CI job (`.github/workflows/ci.yml:296`) run the script. A new test script needs its own invocation, as `check_llama_compat_manifest_test.sh` has at `Makefile:755` and `.github/workflows/ci.yml:373`.
- **Edge cases**: a search directory that contains no matching file at all (the script must fail, not print OK over zero files); a header that contains `cuda_kernel(` inside a comment; a file legitimately removed, which must be distinguishable from a file whose launches moved.
- **Error handling**: a shrunk scope must exit non-zero with a message naming the files that left and how to update the expectation deliberately, in the style of the existing failure message.
## Acceptance Criteria
- [ ] Moving a JIT launch into a header, or into a shared helper outside the token-bearing file, either keeps it in scope or fails the check. Demonstrated by a fixture, not by inspection.
- [ ] The success output distinguishes the number of files scanned from the number actually checked, so a drop in the checked count is visible in CI logs.
- [ ] A companion test covers the shrinking-scope case and is wired into `make verify` and CI, following `scripts/ci/check_llama_compat_manifest_test.sh`.
- [ ] The eight files in scope today are still in scope after the change, and the check still passes on the current tree.
- [ ] `docs/code-guidelines.md:164` describes the new scoping behavior.
## Verification
```
make verify-kernel-dtype-keys
bash scripts/ci/check_kernel_dtype_keys_test.sh
```
A pass is the checker reporting OK on the current tree with both counts printed, and the companion test exiting 0 after asserting that a fixture whose launches moved out of scope makes the checker exit non-zero.
## Technical Considerations
Related: #1053 and #1054 (the bugs this check exists to prevent), #1803 and PR #1869 (the refactor that came close to emptying the scope).
Contributor guide
Research direction
Start with scripts/ci/check_kernel_dtype_keys.py, especially main() and check_file(), then compare the companion-test pattern in scripts/ci/check_llama_compat_manifest_test.sh. Run make verify-kernel-dtype-keys and the stated fixture test while reviewing Makefile:735-738, Makefile:755, and the matching CI entries. Done means moved-launch fixtures fail visibly, scanned and checked counts are distinct, current files still pass, and docs/code-guidelines.md:164 is updated.
Written by the indexing model from the issue text.
Assessment
- Tech stack
- bash, python
- Domain
- ci-cd, documentation, testing
- Issue type
- Refactor
- Difficulty
- 4/5
- Estimated time
- 3-5 days
- Activity status
- Active
- Clarity
- Mostly clear
- Newbie friendliness
- 55/100