refactor(rocke): hoist the shared dispatch selector and kernel-id builder

[AviralGoelAMD]

ISSUE ID : AICK-1513

Two small functions were copy-pasted into every operator family. This moves them into one shared place. No behaviour changes, no kernel code.

The two functions

The pin check. A caller can force a specific kernel instead of letting rocKE route (algorithm="chunk_scan"). Something has to compare that request against each candidate and say yes/no with a reason. "auto", the default, matches anything.

The id builder. Once a kernel is picked it needs an identity — operator, family, candidate, algorithm, arch, ABI, plus hashes of the request and the spec. Logs, tuning records and benchmark rows all refer to a pick by that id.

attention and kda each had their own copy of both, character for character. GDN decode and prefill (PRs 3 and 4 of this stack) would have made copies three and four.

What changed

rocke/dispatch/core.py gains selector_matches() and make_kernel_id(op=...). Both families now call them and keep their old private names as aliases:

_selector_matches = selector_matches                    # both common.py files

def _kernel_id(req, candidate, spec):                   # both __init__.py files
    return make_kernel_id(req, candidate, spec, op="attention")

No call site outside these four files changes.

The one thing worth checking closely

The old builders set family=_FAMILY (each family's own constant). The shared one reads family=candidate.family.

Those are different expressions. family is part of selection_key, the identity tuning records and benchmark rows index by — so if they ever disagreed for some candidate, that candidate's identity would shift silently and old records would stop matching. (It is not part of compile_key, which is arch:abi_version:spec_hash, so nothing recompiles and no wrong binary can be dispatched.)

They agree today: every candidate in both registries is registered with its own family constant. The new test asserts that over the real registries, so it also holds for families added later:

for candidates, family in ((attention_candidates(), ATTENTION_FAMILY),
                           (kda_candidates(), KDA_FAMILY)):
    assert candidates, "registry is empty -- the check would pass vacuously"
    for candidate in candidates:
        assert candidate.family == family

Tests

library/tests/dispatch/test_core_helpers.py is new and tests the helpers directly — the family suites only reached them indirectly, so a bug in one used to surface as a confusing failure somewhere downstream. It covers pin matching (auto, exact pins, both rejection reasons, case/whitespace tolerance), id determinism and spec sensitivity, and the invariant above.

Verification

• library/tests/dispatch → 334 passed, 216 subtests.
• Mutation check: forcing selector_matches to reject everything turns 46 kda dispatch tests red, with the injected reason appearing in the real "no candidate supports request" error. That proves the families route through the shared function rather than a surviving private copy. Reverted, green again.
• black clean. One typing import in kda/common.py, left unused by the change, removed.

Context

PR 2 of a 4-PR stack splitting the GDN gfx950 work (originally one 51-file PR): quad_perm (ROCm#12070) → this → GDN decode → GDN prefill.

It does not depend on ROCm#12070 — that one touches crosslane lowering in platform/, this one touches dispatch. Both sit on develop. ROCm#11795's version of this hoist covered kda only; attention is included here because it held an identical copy, and leaving it would have undercut the point of the change.

Self-reviewed on the fork first: #22.

[Copilot]

🟢 Approval recommended

The refactor preserves existing dispatch identities and selector behavior with focused regression coverage.

Pull request overview

Hoists duplicated dispatch selector and kernel-ID logic into the shared dispatch core without changing kernel behavior.

Changes:

• Adds shared selector_matches and make_kernel_id helpers.
• Routes attention and KDA dispatch through the shared helpers.
• Adds direct helper and registry-invariant tests.

File summaries

File	Description
rocke/dispatch/core.py	Adds shared dispatch helpers.
tests/dispatch/test_core_helpers.py	Tests helper behavior and family consistency.
dispatch/kda/common.py	Reuses the shared selector.
dispatch/kda/__init__.py	Reuses the shared kernel-ID builder.
dispatch/attention/common.py	Reuses the shared selector.
dispatch/attention/__init__.py	Reuses the shared kernel-ID builder.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 0
• Review effort level: Balanced

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.