Expand description
CUDA-graph exec-update (shared, model-agnostic): capture a decode step ONCE, then
re-tune individual kernel nodes’ launch geometry per token via
cuGraphExecKernelNodeSetParams — the llama.cpp graph-serving mechanism (their decode
replays one instantiated graph per token with exact per-token grid shapes; nsys shows
zero launch gaps AND eager-exact grids, where a fixed-bucket replay wastes split blocks).
Mechanism: cuGraphKernelNodeGetParams_v2 returns the node’s CUDA_KERNEL_NODE_PARAMS
whose kernelParams staging is DRIVER-OWNED and stays valid for the node’s lifetime —
scalar args are updated by writing through those pointers, geometry by editing the
struct’s gridDim fields, then cuGraphExecKernelNodeSetParams pushes the new params
into the instantiated exec (topology-preserving update; no re-instantiate).
Safety model: every function here takes the raw handles from a live
cudarc::driver::CudaGraph (which owns destruction); callers must keep that graph
(and the capture keeper) alive while updating/launching.
Structs§
- FaMain
- One fa-decode main node with its paired combine — the per-token geometry-update unit.
- Kernel
Node - One kernel node of a captured graph: raw node handle, its full launch params
(grid/block/smem + driver-owned
kernelParamsstaging), and the resolved symbol name.
Functions§
- fa_
apply - Retune every fa main (and paired combine) in the instantiated exec to the live
t_kv: vec mains get the EAGER split count ns = ceil(t_kv/split_keys(t_kv, nkv)); scalar mains shrink grid.y to their in-kernel ns_eff. No-op when the counts haven’t stepped.split_keysis the caller’s ladder (fa_split_keys) so graph and eager stay in lockstep. - fa_plan
- Classify a captured graph’s fa-decode nodes into per-token-updatable
FaMains. Pairing main->combine is by partO pointer identity (arg staging), not node order. Nodes that aren’t fa mains/combines are left untouched (they replay as captured). - kernel_
nodes - Enumerate every KERNEL node of a captured graph with its launch params and symbol name. Non-kernel nodes (memcpy/memset/empty) are skipped — geometry updates only apply to kernel nodes; everything else replays as captured.
- node_
census - Node-type census of a captured graph (debug: which node types remain — mem-alloc/free nodes are the graph-launch-latency suspects).
- read_
i32_ ⚠arg - Read an i32 scalar argument from the node’s kernelParams staging (see
write_i32_arg). - read_
ptr_ ⚠arg - Read a pointer-valued argument (device pointer as u64) from kernelParams staging.
- set_
exec_ params - Push updated launch params for one node into the instantiated exec.
paramsis the (edited) struct fromkernel_nodes— same node topology, new geometry/arg values. - write_
i32_ ⚠arg - Overwrite one i32 scalar argument in the node’s driver-owned kernelParams staging.
idxis the kernel’s parameter position (launch_builder arg order). The write alone does NOT reach the exec — callset_exec_paramsafter editing to push the change.