chore(ci): kernel dtype-key checker cannot see launches that move into headers

[inureyes]

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).