ROCm/composable_kernel
1
0
Fork 0
mirror of https://github.com/ROCm/composable_kernel.git synced 2026-07-07 15:47:05 +00:00
Code Issues Packages Projects Releases Wiki Activity
3,450 Commits 776 Branches 40 Tags
develop
Commit Graph

5 Commits

Author SHA1 Message Date
Muhammed Emin Ozturk
de292a24f9 [rocm-libraries] ROCm/rocm-libraries#8997 (commit 6e9bfd9) 
feat(ck-tile): TE to dispatcher GEMM bridge (fp16/bf16, all
 layouts) (#8997)
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

> Re-opened from #8479 with a compliant branch name
(users/muozturk/ck-tile/gemm-bridge-all-layout-bf16-fp16). Supersedes
#8479.

## Summary

This PR routes the **Tile Engine (TE) regular-GEMM sweep through the
Dispatcher**,
making the Dispatcher the single source of truth for **codegen → build →
runtime**
while the Tile Engine keeps only the **config search space** and the
**benchmark
loop**. It is the consolidated, **single-commit** GEMM bridge covering
**all four
layouts (`rcr`/`rrr`/`crr`/`ccr`)** and **both `fp16` and `bf16`**.

It is a clean re-roll of the earlier bridge work (previously split
across
#8123 + the stacked key/bf16/layouts/parity/example PRs and consolidated
in
#8261). Those branches accumulated unrelated cross-project commits
through repeated
`develop` merges; **this branch is a single clean commit off the latest
`develop`**
containing only the GEMM-bridge files. It supersedes and replaces #8123
/ #8261.

## Motivation

The Tile Engine historically owned its own codegen/build/runtime for
GEMM
(`tile_engine/ops/gemm/gemm_universal/`). The consolidation goal is for
the
**Dispatcher** to own all of that — exactly as it already does for
**FMHA** and
**Grouped Conv** — so there is one kernel-generation/build/runtime path
and the
TE shrinks to a config+benchmark frontend. This PR brings regular GEMM
in line
with that reference binding.

## The binding (mirrors the FMHA/Conv reference, six stages)

1. **Config JSON (TE side)** — the sweep search space lives in
   `tile_engine/ops/gemm/configs/` (flat op-root layout, matching the
   `fmha/` and `grouped_conv/` bridges).
2. **Codegen (Dispatcher)** —
`dispatcher/codegen/unified_gemm_codegen.py` emits
   one fully-typed `.hpp` per kernel; `GemmKernelConfig.name` reproduces
`KERNEL_NAME` **byte-for-byte** (the thread tying config → kernel →
runtime).
3. **Compile to `.so`** — a single static `gemm_ctypes_lib.cpp` is
force-included
   (`-include <kernel.hpp>`); one `.so` per kernel.
4. **Flat `extern "C"` ABI** — `dispatcher_run_gemm(A, B, C, M, N, K,
time_ms)` +
the kernel-name enumeration entry points. **Host-pointer** memory model
(the C
lib `hipMalloc`s internally) — the FMHA-forward branch of the reference.
5. **Python ctypes wrapper** — `dispatcher/python/gemm_utils.py`
   (`GemmDispatcherLib` + `GpuGemmRunner`).
6. **TE driver (3 phases)** — `gemm_full_benchmark.py` (parallel
codegen+build →
`expand_sweep` → subprocess-isolated benchmark) + the disposable
per-kernel
   worker `run_one_gemm_kernel.py`.

## What's included

**Bridge core**
- `dispatcher/codegen/unified_gemm_codegen.py` — GEMM codegen,
byte-exact naming.
- `dispatcher/bindings/ctypes/gemm_ctypes_lib.cpp` — flat C ABI,
host-pointer model.
- `dispatcher/python/gemm_utils.py` — `GemmKernelConfig`, multi-kernel
build
(`setup_multiple_gemm_dispatchers`), `expand_sweep`,
one-`.so`-per-kernel.
- `tile_engine/ops/gemm/gemm_full_benchmark.py` +
`run_one_gemm_kernel.py` —
  3-phase, multi-GPU, subprocess-isolated driver/worker.

**Feature surface (the point of this PR)**
- **All four layouts** `rcr`/`rrr`/`crr`/`ccr` (row-major C only —
ck_tile rejects
  column-major C at build) with layout-aware host transpose.
- **`fp16` + `bf16`** (bf16 via uint16 byte-encoding; dtype derived from
kernel name).
- **Trait-derived registry `KernelKey`** — replaces the earlier
hard-coded
fp16/rcr key so the registry path generalizes across dtype/layout/tile.

**Correctness & performance hygiene**
- **`--verify`** opt-in fp32 numpy-reference gate (global
`max|out-ref|/max|ref|`),
`verified`/`max_rel` columns in the CSV; a mismatch counts as a failure.
- **Tile Engine AMDGPU `-mllvm` codegen-flag parity** (without these the
kernel
  builds with different occupancy and the timing diverges) and
  **arch-validated tile filtering** against the real pipeline/scheduler.
- **Multi-GPU** fan-out across all visible GPUs (`--devices`,
device-pinned
  `HIP_VISIBLE_DEVICES` workers).

**Example & tests**
- `dispatcher/examples/gemm/python/12_te_bridge.py` — runnable
end-to-end example.
- `dispatcher/tests/test_gemm_parity.py`, `test_gemm_utils.py`, and a
parity
  regression harness.

**Cleanup**
- Removes the legacy standalone `gemm_universal` build path
  (`gemm_universal_instance_builder.py`, `*_benchmark*.{py,cpp,hpp}`,
`gemm_universal/CMakeLists.txt`) and the old
`test/ck_tile/gemm_tile_engine/`
  harness; promotes the sweep configs to the flat op-root `configs/`.

## Design decisions (consistent with the reference)

- **Host-pointer memory ownership** (C lib owns device memory) — matches
FMHA-forward; the Python runner passes host numpy arrays straight
through.
- **One `.so` per kernel** — packaging choice; the multi-kernel name ABI
is
retained (`get_kernel_name_at(0)` reports the single kernel), so the
Python
  enumeration path is unchanged from FMHA/Conv.
- **Flat `configs/`** at the op root — matches the
`fmha/`/`grouped_conv/`
convention; the not-yet-bridged variants keep their per-variant
`configs/`
  dirs, selected by `--variant`.

## Validation (gfx942 / MI300X)

- Bridge build + benchmark + `--verify` across **`fp16` and `bf16`** and
**all
  four layouts**, checked against an fp32 numpy reference (`A @ B`).
- **Name parity** holds end-to-end: each `.so`'s reported runtime name
equals
  `GemmKernelConfig(...).name`.
- bf16 passes under a widened fp16/bf16 tolerance; fp16 within the
standard
  `max_rel` gate.

## Test plan

- [ ] `gemm_full_benchmark.py --verify` over
`configs/default_ci_config.json` for
      `fp16` and `bf16`, each of `rcr`/`rrr`/`crr`/`ccr`.
- [ ] `unified_gemm_codegen.py` emits a header whose stem ==
`GemmKernelConfig.name`.
- [ ] `setup_multiple_gemm_dispatchers` builds + links each config
against
      `gemm_ctypes_lib.cpp`.
- [ ] `pytest dispatcher/tests/test_gemm_parity.py
dispatcher/tests/test_gemm_utils.py`.
- [ ] `examples/gemm/python/12_te_bridge.py` runs end to end.

## Notes

- Single clean commit off the latest `develop`; the diff is **35 files,
all under
`projects/composablekernel/`** (dispatcher + tile_engine/ops/gemm +
test/ck_tile).
- **Supersedes #8123 and #8261**, which will be closed.
- Stream-K (#8136) and grouped GEMM are separate bridge efforts, not in
this PR.
 2026-07-07 01:15:38 +00:00
Vidyasagar Ananthan
b20458e19e [rocm-libraries] ROCm/rocm-libraries#5260 (commit a1834d2) 
[CK] [CK_Tile] Add FMHA scaffolding to CK kernel dispatcher (#5260)

## Motivation

The CK Tile dispatcher currently supports GEMM and Grouped Convolution
but has no support for Fused Multi-Head Attention (FMHA). The
example/ck_tile/01_fmha folder contains a comprehensive FMHA
implementation with forward, backward, split-KV, paged-KV, append-KV,
and batch-prefill kernels across multiple GPU architectures — but there
is no unified dispatch layer for it. This PR ports the FMHA stack into
the dispatcher, following the same architectural patterns established by
GEMM and Grouped Convolution, enabling runtime kernel selection, JIT
compilation from Python, and a declarative C++ example flow. Autotuning
heuristics to follow.

## Technical Details

This PR adds FMHA scaffolding to the CK dispatcher framework, mirroring
GEMM's layered architecture. Seven new C++ runtime headers provide type
definitions (coexisting with upstream headers via __has_include,
requiring zero modifications to example/ck_tile/01_fmha/), a problem
builder with 18+ setters, Signature + Algorithm kernel key matching, a
virtual kernel instance, a DECL_FMHA_KERNEL_SET macro with wildcard
support and named tile/wave/warp setters, arch-aware registry with JSON
export, and a dispatcher with seqtune-aware selection, configurable
timing, and multi-stage execution plans for split-KV (two-stage) and
backward (three-stage). The codegen pipeline is driven by a
fmha_arch_specs.json capturing per-arch tile tables and pipeline
constraints for five architectures (gfx90a/942/950/1100/1201), migrated
from hardcoded logic in 01_fmha/codegen/, with supporting modules for
C++ symbol mappings, validation rules, and named receipt profiles
(ck_default, flash, pytorch, aiter, fp32, fp8). Python integration
(fmha_utils.py) mirrors the C++ layer with JIT compilation, parallel
multi-kernel builds, HIP memory management via ctypes, tolerance-based
validation, and a NumPy CPU reference with GQA support. Twenty-seven C++
and thirty-two Python examples cover the full feature surface — forward,
split-KV, masks, bias, dropout, GQA, backward, append-KV, batch prefill,
fp8, logits soft cap, sink tokens, and parameter sweeps — all
JIT-compiled on the fly.

## Test Plan

Seven test files cover the runtime types, codegen, and end-to-end
correctness. C++ unit tests validate the problem builder, dispatcher
planning (single-stage for forward/paged-KV/append-KV; multi-stage for
split-KV and backward), registry operations, and the kernel-set
declaration macro. Python unit tests verify codegen emission, profile
filtering, and 15 validation rules for masks, hdim constraints, and
pipeline requirements. GPU execution validation in 01_basic_fmha
--validate reports zero errors across 65,536 elements with max absolute
error of 7.29e-05. A gold-standard parity suite (test_fmha_parity.py)
runs 14 configurations through both the upstream tile_example_fmha_fwd
and the dispatcher, comparing exit codes to confirm behavioral parity —
all 14 match.

## Test Result

The C++ smoke test builds and passes all 9 compiled examples, and a
Python JIT sweep (29_sweep_seqlen.py) passes 7/7 configurations reaching
up to 375 TFLOPS at seqlen 2048.

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

---------

Co-authored-by: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com>
Co-authored-by: Mohsen Saffari <mohsen.saffari@amd.com>
Co-authored-by: Maksim (Max) Podkorytov <Maksim.Podkorytov@amd.com>
Co-authored-by: yashagar <yashagar@amd.com>
 2026-05-17 00:29:40 -07:00
Vidyasagar Ananthan
ca28efac88 [rocm-libraries] ROCm/rocm-libraries#5168 (commit 8b5afcb) 
[CK] [CK_Tile] Add GroupConv to Kernel Dispatcher (#5168)

## Motivation

This PR adds CK Tile group convolution (forward, backward-data,
backward-weight) support to the kernel dispatcher, matching and unifying
with the existing dispatcher GEMM infrastructure in architecture and
usability. The dispatcher provides a unified kernel dispatch system with
both C++ and Python frontends, and until now only supported GEMM
operations. This PR enables framework integrators to use the same
declarative kernel workflow for convolutions as they do for GEMM:
declare kernels, build a registry JIT, select kernels within the
registry at runtime, and dispatch to GPU. Future PRs will include
runtime kernel selection heuristics for autotuning of kernel parameters
based on (problem, hardware arch).

## Technical Details

Grouped convolution support has been added to the CK Tile Dispatcher
with generated_conv_backend.hpp enabling dispatcher.run(in, wei, out,
problem) for all 6 conv variants (fwd/bwdd/bwdw x 2D/3D), runtime
heuristic kernel selection, and GroupedConvKernelKey with full
ConvConfigBase fields. Python side adds parallel JIT via
registry.build(max_workers) and heuristic registry.select(). Includes 7
C++ and 6 Python examples covering all directions with CPU reference
validation, and shared infrastructure improvements (BaseRegistry CRTP,
structured exceptions). As a sanity check, JIT compile times for a
single kernel remains the same and for multiple kernels there is better
parallelism:
Kernels | 1 worker | 8 workers
1 | 7.7 s | 7.7 s
2 | 15.9 s | 8.2 s
4 | 33.4 s | 9.7 s
6 | 52.3 s | 10.2 s

## Test Plan

145 ephemeral unit tests have been added to test basic functionality.
All 30 examples/integration tests run end-to-end on gfx950 (MI350): 7
C++ conv, 7 C++ GEMM, 6 Python conv, 10 Python GEMM. CPU reference
validation for forward, backward-data, and backward-weight (2D) in both
C++ and Python examples pass.

## Test Result

30 examples pass. Peak performance: 132 TFLOPS (Batch-32 forward 56x56),
53 TFLOPS (pointwise 1x1). CPU reference accuracy: max_abs_diff < 0.002
for all directions (fp16 vs fp32 reference).

## Submission Checklist

- [x] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

---------

Co-authored-by: Yaswanth Raparti <113389104+yraparti@users.noreply.github.com>
 2026-04-09 10:38:33 -07:00
Yaswanth Raparti
644fc05a87 [rocm-libraries] ROCm/rocm-libraries#5676 (commit 1d18339) 
[CK][CK TILE]Autotuning heuristics infra for universal GEMM kernel selection (#5676)

## Motivation

This PR adds ML-based kernel selection heuristics to the CK Tile
dispatcher, enabling fast and accurate automatic kernel selection for
Universal Gemm kernels. Instead of requiring exhaustive search through
4600+ kernel configurations (taking ~46 seconds per problem shape), the
ML heuristic predicts optimal kernels in microseconds while achieving
>98% of oracle-best performance.

## Technical Details

**ML infrastructure**

https://github.com/ROCm/rocm-libraries/tree/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics
* Feature Engine
([feature_engine.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics/feature_engine.py)):
55-feature extraction including problem dimensions, kernel
configuration, tile efficiency, and hardware profile
* Training Pipeline
([train.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics/train.py)):
LightGBM regression with log-transform, GroupKFold cross-validation,
warm-start support
* Predictor
([predict.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics/predict.py)):
Kernel ranking and TFLOPS prediction for problem shapes
* Evaluation
([evaluate.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics/evaluate.py)):
Comprehensive metrics including efficiency, NDCG@k, shape family
analysis

**Data Generation Tools:**

*
[generate_benchmark_data.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics/generate_benchmark_data.py):
Build and benchmark kernels across diverse problem shapes
*
[convert_json_to_parquet.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics/convert_json_to_parquet.py):
Convert benchmark JSON to training-ready parquet format
*
[data_pipeline.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/heuristics/data_pipeline.py):
Parse streaming benchmark logs into canonical datasets

**Examples**
*
[09_ml_heuristic.cpp](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/examples/gemm/cpp/09_ml_heuristic.cpp):
C++ example demonstrating ML-based kernel selection
*
[09_ml_heuristic.py](https://github.com/ROCm/rocm-libraries/blob/users/vanantha/ck/dispatcher-heuristics/projects/composablekernel/dispatcher/examples/gemm/python/09_ml_heuristic.py):
Python example with validation

**Pre-trained Models
(projects/composablekernel/dispatcher/heuristics/models/):**
* gemm_universal_fp8_gfx950/: fp8 RCR model (42K trees, 97.51% mean
efficiency)
* gemm_universal_fp16_gfx950/: fp16 RCR model (20K trees, 99.36% mean
efficiency)

## Test Plan

* Evaluated on 25 diverse shapes for fp16, 168 shapes for fp8
* All shape families tested: tiny M (M<8), small M, medium M, large M
(M≥1024)
* All pipeline types: compv3, compv4, mem

## Test Result

**fp16 Model (gfx950, RCR layout)**
* Mean Efficiency: 99.36%
* P10 Efficiency: 98.05% (90th percentile of shapes achieve ≥98% of
oracle best)
* Min Efficiency: 95.45%

**fp8 Model (gfx950, RCR layout)**
* Mean Efficiency: 98.28% (original), 97.51% (wide coverage)
* P10 Efficiency: 94.64% (original), 93.89% (wide coverage)
* Min Efficiency: 84.5%

## Submission Checklist

- [x ] Look over the contributing guidelines at
https://github.com/ROCm/ROCm/blob/develop/CONTRIBUTING.md#pull-requests.

---------

Co-authored-by: Vidyasagar Ananthan <vidyasagar.ananthan@amd.com>
Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>
 2026-04-01 19:25:55 -07:00
Vidyasagar Ananthan
9e049a32a1 Adding dispatcher architecture (#3300) 
* WIP POC of dispatcher

* Dispatcher python workflow setup.

* Dispatcher cleanup and updates.

Further dispatcher cleanup and updates.

Build fixes

Improvements and python to CK example

Improvements to readme

* Fixes to python paths

* Cleaning up code

* Improving dispatcher support for different arch

Fixing typos

* Fix formatting errors

* Cleaning up examples

* Improving codegeneration

* Improving and fixing C++ examples

* Adding conv functionality (fwd,bwd,bwdw) and examples.

* Fixes based on feedback.

* Further fixes based on feedback.

* Adding stress test for autogeneration and autocorrection, and fixing preshuffle bug.

* Another round of improvements  based on feedback.

* Trimming out unnecessary code.

* Fixing the multi-D implementation.

* Using gpu verification for gemms and fixing convolutions tflops calculation.

* Fix counter usage issue and arch filtering per ops.

* Adding changelog and other fixes.

* Improve examples and resolve critical bugs.

* Reduce build time for python examples.

* Fixing minor bug.

* Fix compilation error.

* Improve installation instructions for dispatcher.

* Add docker based  installation instructions for dispatcher.

* Fixing arch-based filtering to match tile engine.

* Remove dead code and fix arch filtering.

* Minor bugfix.

* Updates after rebase.

* Trimming code.

* Fix copyright headers.

* Consolidate examples, cut down code.

* Minor fixes.

* Improving python examples.

* Update readmes.

* Remove conv functionality.

* Cleanup following conv removable.
 2026-01-22 09:34:33 -08:00