Skip to main content

Module graph_update

Module graph_update 

Source
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.
KernelNode
One kernel node of a captured graph: raw node handle, its full launch params (grid/block/smem + driver-owned kernelParams staging), 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_keys is 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. params is the (edited) struct from kernel_nodes — same node topology, new geometry/arg values.
write_i32_arg
Overwrite one i32 scalar argument in the node’s driver-owned kernelParams staging. idx is the kernel’s parameter position (launch_builder arg order). The write alone does NOT reach the exec — call set_exec_params after editing to push the change.