Gahow Wang
5b350ee5f0
BF16 greedy decode was sensitive to inter-block scheduling when logits were close, which broke speculative-decoding verify-vs-decode parity. - gemv.cu: write per-K-block partials, then reduce in fixed block order in a second kernel instead of atomicAdd across K-blocks. Scratch buffer size is now n * ceil(k / GEMV_TILE_K); gemv_scratch_elems() exposes this to callers, and decode_graph.rs sizes fp32_hidden/q/kv/ intermediate/vocab from it. - paged_attention.cu: replace atomicAdd merge of warp outputs with per-warp shared partials reduced in warp-id order for both the base and sinks kernels.
650 lines
22 KiB
Rust
650 lines
22 KiB
Rust
//! CUDA Graph integration for batch=1 single-sequence decode.
|
|
//!
|
|
//! Uses a per-layer split graph approach:
|
|
//! - Pre-attention graph: RMSNorm + QKV projections + reshape + QK-norm + RoPE
|
|
//! - Ungraphed: KV cache append + decode attention (variable kv_len)
|
|
//! - Post-attention graph: merge_heads + O-proj + add_rmsnorm + FFN + residual
|
|
//! - Final graph: last RMSNorm + lm_head GEMV
|
|
|
|
use std::ffi::c_void;
|
|
use xserv_cuda::{CudaGraph, CudaStream, GpuBuffer};
|
|
use xserv_kernels::dispatch;
|
|
use xserv_kernels::gemm::{cublas_handle, gemv_scratch_elems};
|
|
|
|
use crate::config::ModelConfig;
|
|
use crate::kv_cache::GpuKVCache;
|
|
|
|
/// Pre-allocated intermediate buffers for decode (batch=1).
|
|
/// All buffers have stable GPU addresses for CUDA Graph replay.
|
|
struct DecodeBuffers {
|
|
// Hidden-size buffers: [1, hidden]
|
|
x: GpuBuffer, // running hidden state
|
|
normed: GpuBuffer, // rmsnorm output
|
|
attn_out: GpuBuffer, // attention output [1, num_heads, 1, head_dim]
|
|
attn_merged: GpuBuffer, // merge_heads output [1, hidden]
|
|
o_proj: GpuBuffer, // O projection output [1, hidden]
|
|
normed2: GpuBuffer, // post-attn norm output [1, hidden]
|
|
sum_out: GpuBuffer, // add_rmsnorm sum output [1, hidden]
|
|
down: GpuBuffer, // down projection output [1, hidden]
|
|
|
|
// QKV projection outputs
|
|
q_proj: GpuBuffer, // [1, num_heads * head_dim]
|
|
k_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
|
|
v_proj: GpuBuffer, // [1, num_kv_heads * head_dim]
|
|
|
|
// Reshaped: [1, H, 1, D]
|
|
q_reshaped: GpuBuffer,
|
|
k_reshaped: GpuBuffer,
|
|
v_reshaped: GpuBuffer,
|
|
|
|
// After QK-norm (same shape as reshaped)
|
|
q_normed: GpuBuffer,
|
|
k_normed: GpuBuffer,
|
|
|
|
// RoPE transposed: [1, H, D]
|
|
q_rope: GpuBuffer,
|
|
k_rope: GpuBuffer,
|
|
|
|
// After RoPE transpose back: [1, H, 1, D]
|
|
q_final: GpuBuffer,
|
|
k_final: GpuBuffer,
|
|
|
|
// FFN intermediates
|
|
gate: GpuBuffer, // [1, intermediate]
|
|
up: GpuBuffer, // [1, intermediate]
|
|
silu_out: GpuBuffer, // [1, intermediate]
|
|
|
|
// GEMV fp32 scratch for deterministic K-block partials.
|
|
fp32_hidden: GpuBuffer, // for hidden-sized GEMV outputs
|
|
fp32_q: GpuBuffer, // for Q projection
|
|
fp32_kv: GpuBuffer, // for K/V projection
|
|
fp32_intermediate: GpuBuffer, // for gate/up projections
|
|
fp32_vocab: GpuBuffer, // for lm_head
|
|
|
|
// Token ID and position (GPU-resident, updated before replay)
|
|
token_id_gpu: GpuBuffer, // 4 bytes (u32)
|
|
position_gpu: GpuBuffer, // 4 bytes (u32)
|
|
|
|
// Final output
|
|
logits: GpuBuffer, // [1, vocab_size]
|
|
}
|
|
|
|
pub struct DecodeGraphState {
|
|
stream: CudaStream,
|
|
buffers: DecodeBuffers,
|
|
|
|
// Per-layer graph pairs
|
|
pre_attn_graphs: Vec<CudaGraph>,
|
|
post_attn_graphs: Vec<CudaGraph>,
|
|
final_graph: CudaGraph,
|
|
|
|
captured: bool,
|
|
|
|
// Model dimensions
|
|
hidden: usize,
|
|
num_heads: usize,
|
|
num_kv_heads: usize,
|
|
head_dim: usize,
|
|
intermediate: usize,
|
|
vocab_size: usize,
|
|
num_layers: usize,
|
|
eps: f32,
|
|
}
|
|
|
|
impl DecodeGraphState {
|
|
pub fn new(config: &ModelConfig) -> Self {
|
|
let hidden = config.hidden();
|
|
let num_heads = config.num_heads();
|
|
let num_kv_heads = config.num_kv_heads();
|
|
let head_dim = config.head_dim();
|
|
let intermediate = config.ffn_hidden();
|
|
let vocab_size = config.vocab_size;
|
|
let num_layers = config.num_layers();
|
|
let eps = config.rms_norm_eps.unwrap_or(1e-6) as f32;
|
|
let es = 2usize; // BF16 = 2 bytes
|
|
|
|
let stream = CudaStream::new().expect("create CUDA stream for graph");
|
|
|
|
let alloc = |size: usize| -> GpuBuffer {
|
|
GpuBuffer::alloc(size).expect("alloc decode graph buffer")
|
|
};
|
|
|
|
let buffers = DecodeBuffers {
|
|
x: alloc(hidden * es),
|
|
normed: alloc(hidden * es),
|
|
attn_out: alloc(num_heads * head_dim * es),
|
|
attn_merged: alloc(hidden * es),
|
|
o_proj: alloc(hidden * es),
|
|
normed2: alloc(hidden * es),
|
|
sum_out: alloc(hidden * es),
|
|
down: alloc(hidden * es),
|
|
|
|
q_proj: alloc(num_heads * head_dim * es),
|
|
k_proj: alloc(num_kv_heads * head_dim * es),
|
|
v_proj: alloc(num_kv_heads * head_dim * es),
|
|
|
|
q_reshaped: alloc(num_heads * head_dim * es),
|
|
k_reshaped: alloc(num_kv_heads * head_dim * es),
|
|
v_reshaped: alloc(num_kv_heads * head_dim * es),
|
|
|
|
q_normed: alloc(num_heads * head_dim * es),
|
|
k_normed: alloc(num_kv_heads * head_dim * es),
|
|
|
|
q_rope: alloc(num_heads * head_dim * es),
|
|
k_rope: alloc(num_kv_heads * head_dim * es),
|
|
|
|
q_final: alloc(num_heads * head_dim * es),
|
|
k_final: alloc(num_kv_heads * head_dim * es),
|
|
|
|
gate: alloc(intermediate * es),
|
|
up: alloc(intermediate * es),
|
|
silu_out: alloc(intermediate * es),
|
|
|
|
fp32_hidden: alloc(
|
|
gemv_scratch_elems(hidden, hidden).max(gemv_scratch_elems(intermediate, hidden))
|
|
* 4,
|
|
),
|
|
fp32_q: alloc(gemv_scratch_elems(hidden, num_heads * head_dim) * 4),
|
|
fp32_kv: alloc(gemv_scratch_elems(hidden, num_kv_heads * head_dim) * 4),
|
|
fp32_intermediate: alloc(gemv_scratch_elems(hidden, intermediate) * 4),
|
|
fp32_vocab: alloc(gemv_scratch_elems(hidden, vocab_size) * 4),
|
|
|
|
token_id_gpu: alloc(4),
|
|
position_gpu: alloc(4),
|
|
|
|
logits: alloc(vocab_size * es),
|
|
};
|
|
|
|
let pre_attn_graphs = (0..num_layers).map(|_| CudaGraph::new()).collect();
|
|
let post_attn_graphs = (0..num_layers).map(|_| CudaGraph::new()).collect();
|
|
|
|
Self {
|
|
stream,
|
|
buffers,
|
|
pre_attn_graphs,
|
|
post_attn_graphs,
|
|
final_graph: CudaGraph::new(),
|
|
captured: false,
|
|
hidden,
|
|
num_heads,
|
|
num_kv_heads,
|
|
head_dim,
|
|
intermediate,
|
|
vocab_size,
|
|
num_layers,
|
|
eps,
|
|
}
|
|
}
|
|
|
|
pub fn is_captured(&self) -> bool {
|
|
self.captured
|
|
}
|
|
|
|
/// Capture all per-layer graphs. Called once after the first decode step.
|
|
pub fn capture(
|
|
&mut self,
|
|
layers: &[LayerWeightPtrs],
|
|
norm_weight: *const c_void,
|
|
lm_head_wt: *const c_void,
|
|
_embed_table: *const c_void,
|
|
rope_cos: *const c_void,
|
|
rope_sin: *const c_void,
|
|
) {
|
|
let s = self.stream.as_raw();
|
|
let h = self.hidden as i32;
|
|
let nh = self.num_heads as i32;
|
|
let nkv = self.num_kv_heads as i32;
|
|
let hd = self.head_dim as i32;
|
|
let inter = self.intermediate as i32;
|
|
let vocab = self.vocab_size as i32;
|
|
let eps = self.eps;
|
|
|
|
let cublas = cublas_handle();
|
|
|
|
// Set cuBLAS to use our stream
|
|
unsafe {
|
|
dispatch::set_cublas_stream(cublas, s);
|
|
}
|
|
|
|
for (l, lw) in layers.iter().enumerate() {
|
|
// === Pre-attention graph ===
|
|
self.pre_attn_graphs[l]
|
|
.begin_capture(&self.stream)
|
|
.expect("begin pre-attn capture");
|
|
unsafe {
|
|
// RMSNorm
|
|
dispatch::rmsnorm_bf16(
|
|
self.buffers.x.as_ptr() as _,
|
|
lw.input_norm,
|
|
self.buffers.normed.as_mut_ptr() as _,
|
|
1,
|
|
h,
|
|
eps,
|
|
s,
|
|
);
|
|
|
|
// Q projection (GEMV)
|
|
dispatch::gemv_bf16(
|
|
self.buffers.normed.as_ptr() as _,
|
|
lw.q_proj_wt,
|
|
self.buffers.q_proj.as_mut_ptr() as _,
|
|
self.buffers.fp32_q.as_mut_ptr() as _,
|
|
h,
|
|
nh * hd,
|
|
s,
|
|
);
|
|
|
|
// K projection (GEMV)
|
|
dispatch::gemv_bf16(
|
|
self.buffers.normed.as_ptr() as _,
|
|
lw.k_proj_wt,
|
|
self.buffers.k_proj.as_mut_ptr() as _,
|
|
self.buffers.fp32_kv.as_mut_ptr() as _,
|
|
h,
|
|
nkv * hd,
|
|
s,
|
|
);
|
|
|
|
// V projection (GEMV)
|
|
dispatch::gemv_bf16(
|
|
self.buffers.normed.as_ptr() as _,
|
|
lw.v_proj_wt,
|
|
self.buffers.v_proj.as_mut_ptr() as _,
|
|
self.buffers.fp32_kv.as_mut_ptr() as _,
|
|
h,
|
|
nkv * hd,
|
|
s,
|
|
);
|
|
|
|
// Reshape heads: [1, H*D] -> [1, H, 1, D]
|
|
dispatch::reshape_heads_bf16(
|
|
self.buffers.q_proj.as_ptr() as _,
|
|
self.buffers.q_reshaped.as_mut_ptr() as _,
|
|
1,
|
|
nh,
|
|
hd,
|
|
s,
|
|
);
|
|
dispatch::reshape_heads_bf16(
|
|
self.buffers.k_proj.as_ptr() as _,
|
|
self.buffers.k_reshaped.as_mut_ptr() as _,
|
|
1,
|
|
nkv,
|
|
hd,
|
|
s,
|
|
);
|
|
dispatch::reshape_heads_bf16(
|
|
self.buffers.v_proj.as_ptr() as _,
|
|
self.buffers.v_reshaped.as_mut_ptr() as _,
|
|
1,
|
|
nkv,
|
|
hd,
|
|
s,
|
|
);
|
|
|
|
// QK norm (head-level rmsnorm: treat [1,H,1,D] as [H, D])
|
|
dispatch::rmsnorm_bf16(
|
|
self.buffers.q_reshaped.as_ptr() as _,
|
|
lw.q_norm,
|
|
self.buffers.q_normed.as_mut_ptr() as _,
|
|
nh,
|
|
hd,
|
|
eps,
|
|
s,
|
|
);
|
|
dispatch::rmsnorm_bf16(
|
|
self.buffers.k_reshaped.as_ptr() as _,
|
|
lw.k_norm,
|
|
self.buffers.k_normed.as_mut_ptr() as _,
|
|
nkv,
|
|
hd,
|
|
eps,
|
|
s,
|
|
);
|
|
|
|
// Transpose for RoPE: [1,H,1,D] -> [1,H,D]
|
|
dispatch::transpose_hsd_to_shd_bf16(
|
|
self.buffers.q_normed.as_ptr() as _,
|
|
self.buffers.q_rope.as_mut_ptr() as _,
|
|
1,
|
|
nh,
|
|
hd,
|
|
s,
|
|
);
|
|
dispatch::transpose_hsd_to_shd_bf16(
|
|
self.buffers.k_normed.as_ptr() as _,
|
|
self.buffers.k_rope.as_mut_ptr() as _,
|
|
1,
|
|
nkv,
|
|
hd,
|
|
s,
|
|
);
|
|
|
|
// RoPE (in-place, reads position_gpu)
|
|
dispatch::rope_bf16(
|
|
self.buffers.q_rope.as_mut_ptr() as _,
|
|
rope_cos,
|
|
rope_sin,
|
|
self.buffers.position_gpu.as_ptr() as _,
|
|
1,
|
|
nh,
|
|
hd,
|
|
s,
|
|
);
|
|
dispatch::rope_bf16(
|
|
self.buffers.k_rope.as_mut_ptr() as _,
|
|
rope_cos,
|
|
rope_sin,
|
|
self.buffers.position_gpu.as_ptr() as _,
|
|
1,
|
|
nkv,
|
|
hd,
|
|
s,
|
|
);
|
|
|
|
// Transpose back: [1,H,D] -> [1,H,1,D]
|
|
dispatch::transpose_shd_to_hsd_bf16(
|
|
self.buffers.q_rope.as_ptr() as _,
|
|
self.buffers.q_final.as_mut_ptr() as _,
|
|
1,
|
|
nh,
|
|
hd,
|
|
s,
|
|
);
|
|
dispatch::transpose_shd_to_hsd_bf16(
|
|
self.buffers.k_rope.as_ptr() as _,
|
|
self.buffers.k_final.as_mut_ptr() as _,
|
|
1,
|
|
nkv,
|
|
hd,
|
|
s,
|
|
);
|
|
}
|
|
self.pre_attn_graphs[l]
|
|
.end_capture(&self.stream)
|
|
.expect("end pre-attn capture");
|
|
|
|
// === Post-attention graph ===
|
|
self.post_attn_graphs[l]
|
|
.begin_capture(&self.stream)
|
|
.expect("begin post-attn capture");
|
|
unsafe {
|
|
// Merge heads: [1,H,1,D] -> [1, hidden]
|
|
// attn_out is written by ungraphed attention
|
|
dispatch::merge_heads_bf16(
|
|
self.buffers.attn_out.as_ptr() as _,
|
|
self.buffers.attn_merged.as_mut_ptr() as _,
|
|
1,
|
|
nh,
|
|
hd,
|
|
s,
|
|
);
|
|
|
|
// O projection
|
|
dispatch::gemv_bf16(
|
|
self.buffers.attn_merged.as_ptr() as _,
|
|
lw.o_proj_wt,
|
|
self.buffers.o_proj.as_mut_ptr() as _,
|
|
self.buffers.fp32_hidden.as_mut_ptr() as _,
|
|
nh * hd,
|
|
h,
|
|
s,
|
|
);
|
|
|
|
// Fused Add+RMSNorm: normed2 = rmsnorm(o_proj + x), sum_out = o_proj + x
|
|
dispatch::add_rmsnorm_bf16(
|
|
self.buffers.o_proj.as_ptr() as _,
|
|
self.buffers.x.as_ptr() as _,
|
|
lw.post_norm,
|
|
self.buffers.normed2.as_mut_ptr() as _,
|
|
self.buffers.sum_out.as_mut_ptr() as _,
|
|
1,
|
|
h,
|
|
eps,
|
|
s,
|
|
);
|
|
|
|
// Gate projection
|
|
dispatch::gemv_bf16(
|
|
self.buffers.normed2.as_ptr() as _,
|
|
lw.gate_proj_wt,
|
|
self.buffers.gate.as_mut_ptr() as _,
|
|
self.buffers.fp32_intermediate.as_mut_ptr() as _,
|
|
h,
|
|
inter,
|
|
s,
|
|
);
|
|
|
|
// Up projection
|
|
dispatch::gemv_bf16(
|
|
self.buffers.normed2.as_ptr() as _,
|
|
lw.up_proj_wt,
|
|
self.buffers.up.as_mut_ptr() as _,
|
|
self.buffers.fp32_intermediate.as_mut_ptr() as _,
|
|
h,
|
|
inter,
|
|
s,
|
|
);
|
|
|
|
// Fused SiLU x Mul
|
|
dispatch::silu_mul_bf16(
|
|
self.buffers.gate.as_ptr() as _,
|
|
self.buffers.up.as_ptr() as _,
|
|
self.buffers.silu_out.as_mut_ptr() as _,
|
|
inter,
|
|
s,
|
|
);
|
|
|
|
// Down projection
|
|
dispatch::gemv_bf16(
|
|
self.buffers.silu_out.as_ptr() as _,
|
|
lw.down_proj_wt,
|
|
self.buffers.down.as_mut_ptr() as _,
|
|
self.buffers.fp32_hidden.as_mut_ptr() as _,
|
|
inter,
|
|
h,
|
|
s,
|
|
);
|
|
|
|
// x = sum_out + down (residual connection for next layer)
|
|
dispatch::add_bf16(
|
|
self.buffers.sum_out.as_ptr() as _,
|
|
self.buffers.down.as_ptr() as _,
|
|
self.buffers.x.as_mut_ptr() as _,
|
|
h,
|
|
s,
|
|
);
|
|
}
|
|
self.post_attn_graphs[l]
|
|
.end_capture(&self.stream)
|
|
.expect("end post-attn capture");
|
|
}
|
|
|
|
// === Final graph: norm + lm_head ===
|
|
self.final_graph
|
|
.begin_capture(&self.stream)
|
|
.expect("begin final capture");
|
|
unsafe {
|
|
dispatch::rmsnorm_bf16(
|
|
self.buffers.x.as_ptr() as _,
|
|
norm_weight,
|
|
self.buffers.normed.as_mut_ptr() as _,
|
|
1,
|
|
h,
|
|
eps,
|
|
s,
|
|
);
|
|
dispatch::gemv_bf16(
|
|
self.buffers.normed.as_ptr() as _,
|
|
lm_head_wt,
|
|
self.buffers.logits.as_mut_ptr() as _,
|
|
self.buffers.fp32_vocab.as_mut_ptr() as _,
|
|
h,
|
|
vocab,
|
|
s,
|
|
);
|
|
}
|
|
self.final_graph
|
|
.end_capture(&self.stream)
|
|
.expect("end final capture");
|
|
|
|
// Reset cuBLAS back to null stream
|
|
unsafe {
|
|
dispatch::set_cublas_stream(cublas, std::ptr::null_mut());
|
|
}
|
|
|
|
self.captured = true;
|
|
}
|
|
|
|
/// Execute a single decode step using captured graphs.
|
|
pub fn execute(
|
|
&mut self,
|
|
token_id: u32,
|
|
position: u32,
|
|
cache: &mut GpuKVCache,
|
|
_layers: &[LayerWeightPtrs],
|
|
embed_table: *const c_void,
|
|
vocab_size: i32,
|
|
hidden_size: i32,
|
|
) {
|
|
assert!(self.captured, "must call capture() before execute()");
|
|
let s = self.stream.as_raw();
|
|
let nkv = self.num_kv_heads;
|
|
let nh = self.num_heads;
|
|
let hd = self.head_dim;
|
|
let es = 2usize; // BF16
|
|
|
|
// Upload token ID and position to fixed GPU buffers
|
|
self.buffers
|
|
.token_id_gpu
|
|
.copy_from_host(&token_id.to_le_bytes())
|
|
.unwrap();
|
|
self.buffers
|
|
.position_gpu
|
|
.copy_from_host(&position.to_le_bytes())
|
|
.unwrap();
|
|
|
|
// Embedding (outside graph since token_id changes each step)
|
|
unsafe {
|
|
dispatch::embedding_bf16(
|
|
embed_table,
|
|
self.buffers.token_id_gpu.as_ptr() as _,
|
|
self.buffers.x.as_mut_ptr() as _,
|
|
1,
|
|
hidden_size,
|
|
vocab_size,
|
|
s,
|
|
);
|
|
}
|
|
|
|
for l in 0..self.num_layers {
|
|
// Pre-attention graph (norm + QKV + reshape + QK-norm + RoPE)
|
|
self.pre_attn_graphs[l]
|
|
.launch(&self.stream)
|
|
.expect("launch pre-attn graph");
|
|
|
|
// Ungraphed: KV cache append
|
|
// k_final shape: [1, num_kv_heads, 1, head_dim] (after RoPE pipeline)
|
|
// v_reshaped shape: [1, num_kv_heads, 1, head_dim] (V skips RoPE)
|
|
let pos = position as usize;
|
|
|
|
let k_buf_size = nkv * hd * es;
|
|
let v_buf_size = nkv * hd * es;
|
|
let shape = [1usize, nkv, 1, hd];
|
|
|
|
// Synchronize before accessing buffers for KV cache append
|
|
self.stream.synchronize().expect("sync before kv cache");
|
|
|
|
let k_view = unsafe {
|
|
crate::kv_cache::tensor_from_gpu_buffer_pub(
|
|
GpuBuffer::borrow_raw(self.buffers.k_final.as_mut_ptr(), k_buf_size),
|
|
&shape,
|
|
xserv_tensor::DType::BF16,
|
|
0,
|
|
)
|
|
};
|
|
let v_view = unsafe {
|
|
crate::kv_cache::tensor_from_gpu_buffer_pub(
|
|
GpuBuffer::borrow_raw(self.buffers.v_reshaped.as_mut_ptr(), v_buf_size),
|
|
&shape,
|
|
xserv_tensor::DType::BF16,
|
|
0,
|
|
)
|
|
};
|
|
cache.append(l, &k_view, &v_view, 1, pos);
|
|
|
|
// Ungraphed: get full KV cache and run decode attention
|
|
let (k_full, v_full) = cache.get_kv_len(l, pos + 1);
|
|
let kv_len = (pos + 1) as i32;
|
|
let scale = 1.0 / (hd as f32).sqrt();
|
|
|
|
// Attention output written to attn_out (separate from q_final)
|
|
unsafe {
|
|
dispatch::decode_attention_bf16(
|
|
self.buffers.q_final.as_ptr() as _,
|
|
k_full.data_ptr() as _,
|
|
v_full.data_ptr() as _,
|
|
self.buffers.attn_out.as_mut_ptr() as _,
|
|
1,
|
|
nh as i32,
|
|
nkv as i32,
|
|
kv_len,
|
|
hd as i32,
|
|
scale,
|
|
s,
|
|
);
|
|
}
|
|
|
|
// Synchronize before post-attention graph reads attn_out
|
|
self.stream.synchronize().expect("sync before post-attn");
|
|
|
|
// Post-attention graph (merge + O-proj + add_rmsnorm + FFN + residual)
|
|
self.post_attn_graphs[l]
|
|
.launch(&self.stream)
|
|
.expect("launch post-attn graph");
|
|
}
|
|
|
|
// Final graph (norm + lm_head)
|
|
self.final_graph
|
|
.launch(&self.stream)
|
|
.expect("launch final graph");
|
|
|
|
// Sync to ensure logits are ready
|
|
self.stream.synchronize().expect("sync after decode");
|
|
}
|
|
|
|
/// Get the logits buffer (for reading results after execute).
|
|
pub fn logits_buffer(&self) -> &GpuBuffer {
|
|
&self.buffers.logits
|
|
}
|
|
|
|
/// Invalidate captured graphs (e.g. when switching sequences).
|
|
pub fn invalidate(&mut self) {
|
|
self.captured = false;
|
|
self.pre_attn_graphs = (0..self.num_layers).map(|_| CudaGraph::new()).collect();
|
|
self.post_attn_graphs = (0..self.num_layers).map(|_| CudaGraph::new()).collect();
|
|
self.final_graph = CudaGraph::new();
|
|
}
|
|
}
|
|
|
|
unsafe impl Send for DecodeGraphState {}
|
|
|
|
/// Lightweight struct holding raw pointers to a layer's weight tensors.
|
|
/// Used to avoid passing the full model struct into the graph capture code.
|
|
pub struct LayerWeightPtrs {
|
|
pub input_norm: *const c_void,
|
|
pub q_proj_wt: *const c_void,
|
|
pub k_proj_wt: *const c_void,
|
|
pub v_proj_wt: *const c_void,
|
|
pub o_proj_wt: *const c_void,
|
|
pub q_norm: *const c_void,
|
|
pub k_norm: *const c_void,
|
|
pub post_norm: *const c_void,
|
|
pub gate_proj_wt: *const c_void,
|
|
pub up_proj_wt: *const c_void,
|
|
pub down_proj_wt: *const c_void,
|
|
}
|
|
|
|
unsafe impl Send for LayerWeightPtrs {}
|
|
unsafe impl Sync for LayerWeightPtrs {}
|