Choosing a dispatch policy

Choosing a dispatch policy#

Pick how a per-sample branch between several formulas is compiled, let the runtime pick it for you, and give the reverse pass its own choice.

Time: ~6 min · Runs on: CPU · You need: hawk and NumPy; read first: Gradients through branches and stops

if / where choose between two values. When a sample picks one of K formulas by an integer class — a regime, a material, a contact state — hawk.math.dispatch(kind, [branch_0, ..., branch_K-1]) says so directly. What it computes never changes; how the branch is compiled is a policy you can set.

Colab or Kaggle: run this first cell (it installs the packages and downloads the helper files this notebook uses). Elsewhere it only checks the imports.

Show code cell source

Hide code cell source

# Colab setup, generated by docs/tools/nb_colab_setup.py (do not edit).
import os

try:
    import google.colab  # noqa: F401
    HOSTED = True  # a Colab runtime
except ImportError:
    HOSTED = "KAGGLE_KERNEL_RUN_TYPE" in os.environ  # a Kaggle kernel

if HOSTED:
    import importlib.metadata, pathlib, re, shutil, subprocess, sys, urllib.request

    gpu = bool(shutil.which("nvidia-smi")) and subprocess.run(["nvidia-smi", "-L"], capture_output=True).returncode == 0
    extra = ""
    if gpu:
        # the CUDA major the runtime already has: a CuPy that is installed, else torch's, else the driver's
        major = None
        for m in ("13", "12"):
            try:
                importlib.metadata.version(f"cupy-cuda{m}x")
                major = m
                break
            except importlib.metadata.PackageNotFoundError:
                pass
        if major is None:
            try:
                import torch
                major = "13" if str(torch.version.cuda).startswith("13") else None
            except ImportError:
                pass
        if major is None:
            smi = re.search(r"CUDA Version: (\d+)", subprocess.run(["nvidia-smi"], capture_output=True, text=True).stdout)
            major = "13" if smi and int(smi.group(1)) >= 13 else "12"
        extra = f"[cuda{major}]"
        print(f"GPU found: installing the {extra} extras to match this runtime.")
    pre = ["--pre"] if os.environ.get("RAPTOR_PRE") else []  # optional: test a release candidate
    ref = os.environ.get("RAPTOR_REF", "main")  # optional: fetch the helper files from this git ref
    if pre or "RAPTOR_REF" in os.environ:
        print(f"Pre-release test: installing with --pre, helper files from {ref}.")
    # our packages with --upgrade, so a runtime that kept an older release gets the current one
    subprocess.check_call([sys.executable, "-m", "pip", "install", "-q", "--upgrade", *pre, f"raptor-hawk{extra}"])
    if not gpu:
        print("Running on the CPU (no GPU in this runtime). For a GPU: " + ("Settings → Accelerator → GPU T4 x2." if "KAGGLE_KERNEL_RUN_TYPE" in os.environ else "Runtime → Change runtime type → T4 GPU."))

try:
    import hawk  # noqa: F401
except ImportError as err:
    raise ImportError(
        "hawk not found. Install with `pip install raptor-hawk`. For a GPU, add [cuda12] (or [cuda13]) to raptor-hawk."
    ) from err

import pathlib, tempfile
import numpy as np
import hawk
import hawk.math as M
from hawk import Index, Mutable, Param, Scalar


@hawk.kernel
def drag(dt: Param, regime: Index, v: Scalar, v_new: Mutable[Scalar]):
    a = M.dispatch(regime, [-0.1 * v,                 # laminar
                            -0.4 * v * abs(v),        # turbulent
                            -2.0 * v * M.exp(-v * v)], # a costly third law
                   policy="auto")
    v_new = v + dt * a


drag

<Kernel drag slots=5>

The four policies#

policy	how it compiles	when it wins
`"switch"` (default)	one kernel, a `switch` on the class	branches of similar cost
`"predicated"`	one kernel, every branch evaluated, one kept	tiny branches, no divergence
`"segmented"`	one kernel per branch, run over the samples of that class	a rare, expensive branch
`"auto"`	compiled as `switch` until a launching runtime chooses	let the runtime pick, from a cost model and the observed branch mix

A policy you declare is never changed by anything else: "auto" is the only value a runtime may resolve, so writing "switch" pins it.

Is the class fixed across steps?#

Splitting a branch into one kernel per class pays for itself when the samples can be sorted by class once and stay sorted. That holds when the class is a static selector: it reads nothing the step writes. hawk records this for every dispatch in the built kernel’s sidecar, together with the planes the selector reads — a runtime intersects those with everything its whole step writes before trusting it:

work_dir = pathlib.Path(tempfile.mkdtemp())
bundle = hawk.build([drag], work_dir, targets=("host",))
(entry,) = bundle.artifacts[0].sidecar["dispatch"]
entry

{'name': 'dispatch0',
 'K': 3,
 'policy': 'auto',
 'static_selector': True,
 'selector_reads': [['per_sample', 'regime']],
 'mix': 'dispatch0_mix'}

regime is an input, so the class is static. A selector computed from the state the kernel updates is not:

@hawk.kernel
def drag_by_speed(dt: Param, cut: Param, slow: Index, fast: Index,
                  v: Mutable[Scalar], regime: Mutable[Index]):
    v = v + dt * M.dispatch(regime, [-0.1 * v, -0.4 * v * abs(v)], policy="auto")
    regime = M.select(abs(v) > cut, fast, slow)   # the step rewrites the class


bundle2 = hawk.build([drag_by_speed], pathlib.Path(tempfile.mkdtemp()), targets=("host",))
(entry2,) = bundle2.artifacts[0].sidecar["dispatch"]
entry2["static_selector"], entry2["selector_reads"]

(False, [['mutable', 'regime']])

A uniform the selector reads (cut above, were it in the selector) is listed too, and still counts as static: the host changes it only between launches, and a runtime that rebinds it re-sorts first.

Counting the branch mix#

To choose between switch and segmented a runtime needs to know how the samples split between the branches. An "auto" dispatch counts it as it runs: one integer counter per branch in the plane the sidecar names (entry["mix"], here dispatch0_mix) — K uint32 cells at the start of the plane, bumped once per live sample (on a GPU, once per warp per branch). The counts are exact integers, the same however the work was scheduled. A dispatch with a declared policy counts nothing and compiles to the same code as before. Inside a kernel that takes several steps per launch the count is per step taken, so the mix weighs every step alike. Left unbound, the counts go to a scratch plane; bind your own to read them:

n = 9
regime = np.array([0, 0, 1, 2, 2, 2, 1, 0, 5], dtype=np.int64)   # 5 clamps to 2
mix = np.zeros(3, dtype=np.uint32)       # one counter per branch
hawk.run(hawk.load(work_dir, "drag"), dt=0.1, regime=regime, v=np.ones(n),
         v_new=np.zeros(n), dispatch0_mix=mix)
print("branch counts:", mix, " bincount:",
      np.bincount(np.clip(regime, 0, 2), minlength=3))

branch counts: [3 2 4]  bincount: [3 2 4]

hawk.ir.dispatch_costs(walk) is the cost view a runtime reads to make that choice. It takes the host profile (host_profile=, default the one in effect) and reports, per dispatch group, whether the fused host loop is fused_vectorised: true for a sample-local kernel whose branches call no library math, or whose host profile vectorises it (native-vector-math). Under an exact profile (native, portable) a call like exp stays scalar, so the fused loop is not counted as running every branch in lockstep.

Resolving `"auto"`#

hawk.resolve_dispatch(kernel, {group: policy}) is how the choice is applied: the same kernel, with that group lowered as the policy named. Only an "auto" group can be resolved; one you declared is refused, so your choice always stands. A kernel that steps several times per launch keeps doing so under switch or predicated; resolved to segmented it becomes its single step, split per branch, one step per launch.

drag_split = hawk.resolve_dispatch(drag, {"dispatch0": "segmented"})
split_bundle = hawk.build([drag_split], pathlib.Path(tempfile.mkdtemp()), targets=("host",))
print([a.name for a in split_bundle.artifacts])

['drag__k0', 'drag__k1', 'drag__k2', 'drag__cls']

Each unit runs over the samples of its own class, which the runtime keeps sorted. If the class of a sample changes after that sort — a selector the step writes — the unit does not apply the wrong formula: it re-reads the selector, keeps the sample’s old value, marks it in dispatch0_stale and writes its new class to dispatch0_class; the runtime finishes the stale samples with the switch form and re-sorts.

The reverse pass takes its own structure#

vjp derives a dispatch per dispatch: by default each derived dispatch inherits the forward’s policy. The reverse pass does different arithmetic, so the best structure can differ; policy={group: ...} sets it, naming the group the way the sidecar does (dispatch0, dispatch1, … in order) — for example vjp(drag, wrt=("v",), policy={"dispatch0": "segmented"}). Here the inherited policy, checked against a finite difference:

from hawk import Kernel
from hawk.diff import vjp

grad_switch = Kernel("drag_vjp", vjp(drag, wrt=("v",)))
bundle3 = hawk.build([drag, grad_switch], pathlib.Path(tempfile.mkdtemp()), targets=("host",))

n = 6
regime = np.array([0, 1, 2, 0, 1, 2], dtype=np.int64)
v = np.linspace(-1.0, 1.0, n)
bar_v = np.zeros(n)
hawk.run(hawk.load(bundle3.directory, "drag_vjp"), dt=0.1, regime=regime, v=v,
         bar_v_new=np.ones(n), bar_v=bar_v)

h = 1e-6
def step(v):
    out = np.zeros(n)
    hawk.run(hawk.load(bundle3.directory, "drag"), dt=0.1, regime=regime, v=v, v_new=out)
    return out

print("vjp          :", np.round(bar_v, 8))
print("finite diff  :", np.round((step(v + h) - step(v - h)) / (2 * h), 8))

vjp          : [0.99       0.952      0.82321474 0.99       0.952      1.07357589]
finite diff  : [0.99       0.952      0.82321474 0.99       0.952      1.07357589]

The value of the gradient does not depend on the policy — only how it is compiled. An unknown group or policy is refused by name:

try:
    vjp(drag, wrt=("v",), policy={"dispatch3": "switch"})
except hawk.HawkError as err:
    print(err)

vjp(policy=...): 'dispatch3' is not a dispatch group of drag; its groups are ['dispatch0']

What just happened#

dispatch picks one of K formulas per sample; the policy only changes how it is compiled.
"auto" hands the choice to the launching runtime and counts the branch mix it needs; hawk.resolve_dispatch applies the choice. Every other policy is yours and stays yours.
The sidecar records whether each selector is static and which planes it reads, so a runtime can decide whether sorting by class pays.
vjp(..., policy={group: ...}) lets the reverse pass pick its own structure; the gradient is the same either way.

Try this#

Resolve drag to "predicated" instead, build it, and compare the source with the switch form: every branch is evaluated and one kept.

Next#

Writing a vocabulary — declare a reusable kernel shape once, so a whole family of models shares one vocabulary.