TheRock/rocm-libraries/dnn-providers/hip-kernel-provider/rocke/library/dispatch/AGENTS.md

# Attention dispatcher — agent guide

This folder owns the **path-level dispatch** for unified attention: which kernel
family (2D tiled prefill vs 3D split-KV decode) and which specialized candidate
handles a given `AttentionRequest`.

## What the dispatcher decides (and what it does NOT)

**Decides:** kernel path (`"2d"` or `"3d"`), candidate name, algorithm tag,
and spec identity `(path, head_size, block_size)`.

**Does NOT decide:** CTA geometry (`num_warps`, `tile_size`), `num_segments`,
`waves_per_eu`, or any other performance knob. Those live in
`builders/common/attention_spec_builder.py` and
`kernels/common/attention_unified.py`. The parity identity with C++ is
`(path, head_size, block_size)` only — C++ reads `num_segments` as a parameter
passed from Python, it does not recompute it.

## Candidate registry — priority table

| priority | candidate | scope |
|---|---|---|
| 5 | `attention_gfx942_dense_pipe` | gfx942 fp16 2D prefill flash |
| 5 | `attention_d256_decode` | gfx942/gfx950 bf16 D256 3D decode |
| 10 | `attention_unified_2d` | generic 2D prefill fallback |
| 10 | `attention_unified_3d` | generic 3D decode fallback |

Lower priority number = higher precedence. Generic candidates (10) remain the
fallback for everything a specialized candidate (5) does not claim.

## How to add a new specialized candidate

Follow the `_make_d256_decode_candidate()` pattern in `attention.py`:

1. **Add a cohort predicate** in `kernels/common/attention_unified.py` —
   a pure function of `UnifiedAttentionProblem` that returns `True` for the
   target shape family. This is the single source of truth for membership;
   import it lazily inside the factory to keep `dispatch/` arch-neutral.

2. **Add a factory function** `_make_<name>_candidate()` in `attention.py`.
   The `support()` closure must check in order: request errors → arch/dtype
   gates → `_selector_matches` → `supports_native_unified_attention` →
   cohort predicate → path check (`select_path() == "2d"` or `"3d"`).

3. **Register** with `ATTENTION_REGISTRY.register(_make_<name>_candidate())`.

4. **Add CPU-only dispatch tests** in
   `tests/dispatch/attention/test_<name>_wiring.py`. Cover: registration,
   spec_id, algorithm, priority ordering, support gates (wrong arch/dtype/
   cohort/path), and routing for each target arch. Use `_PinnedArch` context
   manager to avoid GPU dependency.

5. **No C++ changes needed.** The dispatcher is Python-only.

## How to tune `num_segments` for a new cohort

See the worked example:
`benchmarks/gfx942/attention/decode/TUNING_D256.md`

## Testing (CPU-only, no GPU)

```bash
PYTHONPATH=library:platform/python python -m unittest discover \
    -s library/tests/dispatch/attention -p "test_*.py" -v
```

All dispatch tests complete in < 1 s.