Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions .gitignore
Original file line number Diff line number Diff line change
Expand Up @@ -10,3 +10,7 @@ tools/kernel-capture/libkernelcapture.so
/.wrangler/
/attestation.json
/kernels-qwen38-dflash2
kernels-c9/
kernels-gdnT/
kernels-varA/
kernels-varB/
8 changes: 8 additions & 0 deletions crates/kern-manifest/src/types.rs
Original file line number Diff line number Diff line change
Expand Up @@ -586,6 +586,14 @@ pub struct Step {
/// Dynamic shared memory in bytes, if the step needs any.
#[serde(default, skip_serializing_if = "Option::is_none")]
pub shared_mem: Option<Expr>,
/// Programmatic dependent launch: the runtime launches this step with
/// `CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION`, so it may
/// start while the preceding launch is still draining. Only for kernels
/// that put every read of an upstream product and every global write
/// behind `griddepcontrol.wait` (their own inputs — weights, state —
/// may stream ahead of it).
#[serde(default, skip_serializing_if = "is_false")]
pub pdl: bool,
/// Wiring: where each step param comes from — a forwarded interface
/// arg, a scratch buffer, or an implementation-private literal.
pub args: Vec<StepArg>,
Expand Down
3 changes: 3 additions & 0 deletions crates/kern-runtime/src/compile.rs
Original file line number Diff line number Diff line change
Expand Up @@ -68,6 +68,8 @@ pub(crate) enum LaunchKind {
block: [u32; 3],
grid: [CExpr; 3],
shared_mem: Option<CExpr>,
/// Launch with programmatic stream serialization (see `Step::pdl`).
pdl: bool,
},
/// `extern:cublaslt_bf16_tn` / `..._acc` (beta 0.0 / 1.0).
Gemm { beta: f32 },
Expand Down Expand Up @@ -350,6 +352,7 @@ fn compile_dispatch(
compile_expr(&st.grid[2], syms)?,
],
shared_mem: st.shared_mem.as_ref().map(|e| compile_expr(e, syms)).transpose()?,
pdl: st.pdl,
},
};
launches.push(Launch {
Expand Down
30 changes: 29 additions & 1 deletion crates/kern-runtime/src/lib.rs
Original file line number Diff line number Diff line change
Expand Up @@ -626,7 +626,7 @@ impl Runtime {
}
match &l.kind {
LaunchKind::Gemm { beta } => gemm_bf16_tn(&self.blt, &self.stream, &vals, *beta),
LaunchKind::Cubin { func, block, grid, shared_mem } => {
LaunchKind::Cubin { func, block, grid, shared_mem, pdl } => {
let grid =
(grid[0].eval(env)? as u32, grid[1].eval(env)? as u32, grid[2].eval(env)? as u32);
let smem = match shared_mem {
Expand All @@ -638,6 +638,34 @@ impl Runtime {
let raw: Vec<u64> = vals.iter().map(|r| r.val).collect();
let mut params: Vec<*mut c_void> =
raw.iter().map(|s| s as *const u64 as *mut c_void).collect();
if *pdl {
// Programmatic dependent launch: inside stream capture
// this becomes a programmatic graph edge, so the kernel
// may begin (and stream its own inputs) while the
// previous launch drains; its griddepcontrol.wait keeps
// the data dependency.
let mut attr = sys::CUlaunchAttribute {
id: sys::CUlaunchAttributeID::CU_LAUNCH_ATTRIBUTE_PROGRAMMATIC_STREAM_SERIALIZATION,
pad: [0; 4],
value: sys::CUlaunchAttributeValue { programmaticStreamSerializationAllowed: 1 },
};
let cfg = sys::CUlaunchConfig {
gridDimX: grid.0,
gridDimY: grid.1,
gridDimZ: grid.2,
blockDimX: block[0],
blockDimY: block[1],
blockDimZ: block[2],
sharedMemBytes: smem,
hStream: self.stream.cu_stream(),
attrs: &mut attr,
numAttrs: 1,
};
let r = unsafe {
sys::cuLaunchKernelEx(&cfg, *func, params.as_mut_ptr(), std::ptr::null_mut())
};
return cuda_check(r, "cuLaunchKernelEx");
}
unsafe {
cu::launch_kernel(
*func,
Expand Down
Loading
Loading