Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
26 changes: 26 additions & 0 deletions .github/workflows/ci.yml
Original file line number Diff line number Diff line change
Expand Up @@ -364,6 +364,32 @@ jobs:
- name: Check every CUDA JIT launch keys its cache on the input dtypes
run: python3 scripts/ci/check_kernel_dtype_keys.py

kernel-port-dispatch:
name: kernel port dispatch
# Not behind the `changes` filter, for the same reasons as the job above: no
# toolchain, seconds to run, and the defect it guards against is silent.
#
# The defect is specifically a convention decaying. Choosing a kernel's port
# with `use_cuda ? cuda : metal` reads "not CUDA" as "Metal", which was true
# with two backends and wrong with three: on ROCm the false arm threw and,
# because the bridge declarations were not `Result`, the throw crossed a
# `noexcept` extern into `std::terminate`. #1803 fixed the idiom and left
# nine sites shaped that way; #1885 and #2018 then guarded them by hand, one
# launcher at a time, months apart, each after it had already aborted a gate
# run, and the hand-written guards drifted in wording and twice named a
# predicate that did not exist. A convention that must be re-applied at every
# new launcher gets missed at some of them, and only the backend that lacks
# the port ever finds out. This job is what keeps the single helper single.
runs-on: ubuntu-latest
permissions:
contents: read
steps:
- uses: actions/checkout@v7
with:
persist-credentials: false
- name: Check every fused-kernel launcher routes through select_kernel_port
run: python3 scripts/ci/check_kernel_port_dispatch.py

binary-assets:
name: binary assets
# Deliberately not behind the `changes` filter, for the same reason as
Expand Down
9 changes: 7 additions & 2 deletions Makefile
Original file line number Diff line number Diff line change
Expand Up @@ -801,7 +801,7 @@ verify-rocm-smoke: ## ROCm smoke: build, link and generate on the GPU, asserting
# "does a real forward pass still run on this device" answer even while the
# suite is red.
.PHONY: verify-rocm
verify-rocm: verify-versions verify-kernel-dtype-keys verify-llama-compat verify-fmt verify-clippy-rocm verify-rocm-smoke verify-test-rocm ## Run the ROCm gate locally on an AMD host (issue #1811)
verify-rocm: verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt verify-clippy-rocm verify-rocm-smoke verify-test-rocm ## Run the ROCm gate locally on an AMD host (issue #1811)
@echo "$(GREEN)[verify-rocm] OK$(RESET)"

.PHONY: verify-versions
Expand All @@ -814,6 +814,11 @@ verify-kernel-dtype-keys: ## Assert every CUDA JIT kernel launch keys its cache
@echo "$(CYAN)[verify] kernel dtype cache keys...$(RESET)"
@python3 scripts/ci/check_kernel_dtype_keys.py

.PHONY: verify-kernel-port-dispatch
verify-kernel-port-dispatch: ## Assert every fused-kernel launcher chooses its port through select_kernel_port (issues #1801, #1885)
@echo "$(CYAN)[verify] kernel port dispatch...$(RESET)"
@python3 scripts/ci/check_kernel_port_dispatch.py

.PHONY: verify-binary-assets
verify-binary-assets: ## Assert every tracked binary file is declared with a size budget
@echo "$(CYAN)[verify] tracked binary assets...$(RESET)"
Expand Down Expand Up @@ -959,7 +964,7 @@ bump-version: ## Release: set every version-tracking crate to VERSION and sync C
@$(MAKE) --no-print-directory verify-versions

.PHONY: verify
verify: verify-versions verify-kernel-dtype-keys verify-llama-compat verify-fmt verify-clippy verify-test ## Run the full CI-faithful gate locally (recommended before push)
verify: verify-versions verify-kernel-dtype-keys verify-kernel-port-dispatch verify-llama-compat verify-fmt verify-clippy verify-test ## Run the full CI-faithful gate locally (recommended before push)
@echo "$(GREEN)[verify] OK: matches the nightly-verify GitHub Actions job$(RESET)"

.PHONY: verify-clean
Expand Down
Original file line number Diff line number Diff line change
@@ -0,0 +1,76 @@
# Technical Report: PR #2026, choose custom-kernel ports through one helper

**Date**: 2026-09-29

**Status**: Implemented and validated on the gfx1151 host; pending merge. Metal and CUDA were not exercised, since neither is available here; the argument that their behavior is unchanged is by construction and is stated below.

**Languages**: C++ (kernel dispatch), Python (CI checker), Make, GitHub Actions

**Risk level**: Medium. Nine launchers change how they obtain a kernel, though the resolved kernel is identical on every backend that has a port.

## Executive summary

Every fused kernel chose its per-backend port by hand, as `use_cuda ? cuda_port : metal_port`. That expression reads "not CUDA" as "Metal". It was correct while Metal and CUDA were the only backends and became a defect the moment a third one existed. This PR replaces all nine hand-written choices with one helper, adds a table per kernel that both the dispatch and the support predicate read, and adds two mechanical defenses so the pattern cannot regrow.

The change fixes no user-visible bug on its own. What it removes is the need to remember something at every future launcher and every future backend.

## Problem statement

The defect this pattern produced is documented across three issues, and the sequence is the argument for standardizing rather than continuing to patch.

Issue #1803 found the original idiom, `!metal::is_available()` read as "CUDA", and replaced it with a named backend kind. It fixed the CUDA arm and left the false arm still meaning Metal, at nine sites. On ROCm those sites took the Metal arm, `fast::metal_kernel` threw, and because the bridge declarations were not `Result` the throw crossed a `noexcept` cxx extern and ended the process rather than failing the call.

Issue #1885 then fixed two of the nine, the samplers, after they aborted a gate run. PR #2018 fixed the other four two-way launchers plus the paged ones, also after a gate run surfaced them, months later. Each fix was correct in isolation and the set of them drifted: two guard comments cited support predicates that do not exist in the tree, one hardcoded an entry-point name that is wrong for one of its two callers, and the refusal wording differed per site.

A separate inconsistency had the same root. `ssm_kernel_available()` answered the "is there a port" question with `cu::is_available()` while the launcher it gates guarded on `custom_kernels_available()`. Two predicates for one question, with nothing keeping them in agreement.

None of this is a story about careless edits. It is what happens to a convention that has to be re-applied at every new launcher: it gets missed at some of them, and only the backend that lacks the port ever finds out, at runtime, as an abort.

## Change summary

- `src/lib/mlx-cpp/turbo/kernel_port.h` and `.cpp` (new): a `KernelPorts` table of per-backend getters, `select_kernel_port` which resolves one or throws a message naming the entry point and the fallback, and `has_kernel_port` which answers the support predicate from the same table.
- `src/lib/mlx-cpp/turbo/gpu_backend.h`: `kGpuKernelBackendCount`, pinned by a `static_assert` in `kernel_port.cpp`.
- Nine launchers converted: `gumbel_max_sample`, `rejection_sample`, `fused_add_rms_norm`, `fused_rope_qk_append`, `ssm_update_kernel`, `run_fused_moe_two_kernel`, `paged_attention_decode`, `paged_attention_decode_v2_partial`, `paged_attention_merge_states`. No backend comparison remains outside the helper.
- `scripts/ci/check_kernel_port_dispatch.py` (new), a `verify-kernel-port-dispatch` make target added to both `verify` and `verify-rocm`, and an unconditional hosted CI job.
- `run_fused_moe_two_kernel` takes its caller's entry-point name, fixing a mislabel introduced by PR #2018.

## Technical decisions

### The enforcement is a compile error, not a warning

`-Wswitch` already flags a new enumerator in an exhaustive switch, and relying on that was considered and rejected. This repository compiles C++ with warnings enabled but not as errors, and the MLX build emits dozens of them on every run, so a new warning would not be seen. The `static_assert` on `kGpuKernelBackendCount` is a hard stop whose message names the three places to extend. Adding a backend now cannot silently read as "no port" at every kernel.

### All nine, not a subset

Converting some launchers and leaving others would mean the checker has to exempt the remainder, and an exemption list is the mechanism by which the old pattern would survive. Exemptions that do exist are named with a reason and a tracking issue rather than left as bare entries.

### A null table entry rather than an absent case

"No port for this backend" is written into the table as `nullptr` instead of being what a dispatch falls through to. That is the whole difference between this shape and the old one. It also means a Metal-only kernel and a fully ported one have the same shape, and it gives issue #1814 a named slot to fill rather than a dispatch to restructure.

### The support predicate reads the dispatch's table

This is the part that removes an existing bug class rather than a hypothetical one. With `has_kernel_port`, a predicate and its launcher cannot disagree, which makes `expect` at a gated call site correct by construction rather than by inspection.

### Two launchers left out deliberately

`turbo4_delegated_sdpa.cpp` and `sparse_v_sdpa.cpp` are Metal-only and have no refusal at all today. Adding one is a behavior change rather than a mechanical edit: `sparse_v_available()` gates on `turbo::sparse_v::is_enabled()`, an environment threshold, and `KVCacheMode::Turbo4Asym`, with no backend or port term, so whether the launcher is reachable off Apple is unresolved. Both are listed in the checker with that reason.

## Validation

- Full ROCm gate on gfx1151: one failing target, `-p mlxcel-core --lib`, unchanged from before this branch. That failure is the nvfp4 abort of issue #1806 and the only `terminate called` in the run.
- `cargo fmt --all -- --check`, `cargo clippy -p mlxcel --lib --tests -- -D warnings` (the command CI runs), and `actionlint` with shellcheck at error severity: all clean.
- Each of the three defenses was confirmed by negative control rather than assumed to work. Reintroducing a two-way dispatch and hand-rolling a refusal each produce a checker failure naming the file and line. Adding a `Vulkan` enumerator fails the build on the `static_assert` with the intended message.

### Method note worth keeping

The bulk conversion was scripted, and the script placed a table inside a function body in three files because its insertion anchor was "the preceding blank line", which is not a structural boundary in C++. Two of those were caught immediately; the third survived a verification step that counted brace depth, which cannot distinguish a namespace from a function. The compiler caught all three. When editing C++ structure textually, the anchor has to be structural (here, the closing brace of a named holder accessor) and the verification has to be a compile, not a heuristic.

## What is not verified

- **Metal and CUDA were not run.** Neither is available on this host. The argument that their behavior is unchanged is that for any backend with a port the resolved kernel is the same object the old expression resolved, and the new throw path requires the table entry to be null, which it is not on those backends.
- **The refusal path was not re-probed after the refactor.** It was verified by execution before this change, in PR #2018, and the message text moved into the helper unchanged in shape. The gate exercises the success path on ROCm only through the launchers that fall back earlier.

## Remaining work

Issue #1814 fills the `rocm` entries these tables now carry. Two follow-ups fall out of this work: bringing the two Metal-only SDPA launchers into the table once their reachability off Apple is settled, and resolving whether `sparse_v_available()` needs a backend term, which is a correctness question rather than a performance one.
Loading
Loading