mirror of
https://github.com/leejet/stable-diffusion.cpp.git
synced 2026-11-12 11:56:36 +00:00
refactor: register RoPE kernels and fusion outside ggml
This commit is contained in:
1 parent
02e664d12d
commit
f4a1902061
7 files changed
+336
-5
No files matched your search
@@ -0,0 +1,107 @@
|
||||
# ggml operator extensions
|
||||
|
||||
The patched ggml exposes `ggml-ext.h` for application-owned operators, kernel
|
||||
overrides, and graph fusion rules. Registrations are in-process and scoped to a
|
||||
backend device. Register them before capability queries or graph allocation;
|
||||
callback code and user data must remain alive until all dependent graphs and
|
||||
backends have been destroyed. Registration names are unique, and duplicate
|
||||
registrations fail rather than replacing live implementations.
|
||||
|
||||
## Operator definitions
|
||||
|
||||
`ggml_ext_register_op` registers a versioned name, attribute size, and an output
|
||||
metadata inference callback. `ggml_ext_op_create` creates an ordinary ggml tensor
|
||||
with explicit input dependencies. The initial interface supports one contiguous,
|
||||
out-of-place output with an existing `ggml_type`.
|
||||
|
||||
Attributes are copied into the unused portion of `op_params`. They are small,
|
||||
plain values; use `memcpy` when reading attributes that need stricter alignment.
|
||||
Auxiliary tensors must appear in `src`, not in an attribute pointer or user data.
|
||||
Registered operators use `GGML_OP_CUSTOM` internally, with an identifiable
|
||||
callback marker. Unsupported devices do not mistake them for CPU callbacks.
|
||||
These process-local descriptors are not a graph serialization format.
|
||||
|
||||
## Kernel selection
|
||||
|
||||
`ggml_ext_register_kernel(device, &desc)` associates a kernel with a built-in
|
||||
`ggml_op`, or with a registered custom operator. `supports` describes correctness
|
||||
constraints; `preferred` optionally decides whether to override an available
|
||||
native implementation. Higher priorities are considered first. When native
|
||||
execution is unsupported, a supported extension can run regardless of its
|
||||
preference callback.
|
||||
|
||||
Capability queries include registered single-node implementations, so the
|
||||
scheduler can select a device even when its native implementation rejects the
|
||||
operator. Callbacks inspect metadata only and must not recursively call the
|
||||
public capability query for the same device and operator. A fusion-only
|
||||
registration does not claim single-node support.
|
||||
|
||||
`prepare` creates optional host-side plan state and reports workspace bytes.
|
||||
It can run during memory measurement, before tensor addresses are assigned, and
|
||||
must not launch work or allocate device buffers. Workspace is allocated by ggml
|
||||
before execution and included in scheduler memory estimates and buffer-size
|
||||
reports. `execute` receives the backend, workspace buffer, and current tensor
|
||||
addresses. `destroy` releases prepared state. A preparation or execution error
|
||||
is propagated; execution never retries a native implementation after launch.
|
||||
|
||||
Overriding a built-in operator preserves its numerical and aliasing contract,
|
||||
including allocator reuse of input storage. Explicit views/in-place nodes are
|
||||
not intercepted by this first interface. Rebuild the scheduler allocation after
|
||||
changing graph metadata, selection policy, or registrations.
|
||||
|
||||
## Fusion
|
||||
|
||||
A `match` callback returns the number of consecutive nodes it can replace.
|
||||
The scheduler freezes this selection before allocating tensors. ggml checks
|
||||
compute flags, external uses, output flags, and views; the initial implementation
|
||||
accepts a conservative set of pure operators with only the final node exposed
|
||||
as the fused output. The rule remains responsible for mathematical equivalence,
|
||||
dtype, shape, layout, and hardware constraints.
|
||||
|
||||
External inputs stay allocated through the final node. Native graph optimization
|
||||
and execution operate on the remaining contiguous ranges, so native fusion cannot
|
||||
consume nodes selected by an extension. Intermediate allocations are retained
|
||||
in this initial implementation; only their computation is removed.
|
||||
|
||||
Automatic extension fusion is enabled in scheduler execution with one pipeline
|
||||
copy and without evaluation callbacks. Multiple pipeline copies and evaluation
|
||||
callbacks retain individual operations. Direct backend graph computation supports
|
||||
single-node extensions; it does not fuse an already allocated graph. Direct calls
|
||||
use temporary plans and synchronize before releasing them; scheduler plans and
|
||||
workspaces are reused across executions.
|
||||
|
||||
Use the backend graph APIs or the scheduler. The legacy CPU `ggml_graph_compute`
|
||||
and backend `graph_plan_*` interfaces do not dispatch these extensions.
|
||||
|
||||
## Backend services
|
||||
|
||||
Services are discovered through `ggml_backend_reg_get_proc_address`:
|
||||
|
||||
| Service | Purpose |
|
||||
| --- | --- |
|
||||
| `ggml_ext_cpu_parallel_v1` | Run an application callback on the CPU backend's configured thread pool, with its abort callback. |
|
||||
| `ggml_ext_cuda_context_v1` | Borrow the current CUDA device and stream, ordered with native work. |
|
||||
| `ggml_ext_cuda_device_v1` | Get the physical CUDA ordinal for capability checks before backend creation, including virtual ggml devices. |
|
||||
|
||||
Workspace is a backend buffer. Its base is a CUDA device pointer on CUDA; callers
|
||||
must use the representation appropriate to their backend. A Vulkan command
|
||||
recording service is not implemented yet.
|
||||
|
||||
## sd.cpp integration
|
||||
|
||||
`src/core/ggml_extensions.cpp` registers the CPU implementation of
|
||||
`sd.rope_apply.v1`, an override for existing `GGML_OP_ROPE_APPLY` nodes, and a
|
||||
`RMSNorm -> Mul -> RoPE` fusion. The fusion accepts a one-dimensional weight
|
||||
vector. Both interleaved and split-half RoPE, partial rotation, and strided input
|
||||
rows are supported. Model construction uses the custom operator on CPU and
|
||||
retains existing backend implementations elsewhere.
|
||||
|
||||
`SD_USE_UPSTREAM_GGML` excludes these registrations. The stable-diffusion public C
|
||||
API is unchanged. INT8 ConvRot retains its existing packed activation layout and
|
||||
separate weight scale tensor.
|
||||
|
||||
CuTe Flash Attention still resides in ggml, including its existing fusion paths.
|
||||
Migrating it requires moving its pool/conversion dependencies onto the extension
|
||||
workspace and CUDA service, then registering its selection and fusion rules in
|
||||
sd.cpp. This initial implementation establishes and validates the mechanism;
|
||||
it does not yet externalize all model-specific kernels or add dynamic dtypes.
|
||||
Reference in new issue
Block a user