ops: embedding/reshape/transpose/split-merge-heads fwd+bwd
Phase T5 structural ops on top of the T4 set, needed to assemble the tiny transformer: - embedding: gather rows by I32 ids (CUDA kernel) / scatter-add backward (atomic, so repeated ids accumulate). csrc/ops/model.cu + ffi. - reshape: contiguous metadata-only view (Tensor::reshape), no kernel. - transpose_3d01: [a,b,c]->[b,a,c] for the multi-head layout (kernel). - autograd nodes: embedding/reshape/transpose_3d01/transpose_2d, plus split_heads (->Vec<Var>) / merge_heads for per-head attention. - tape: Var::zero_grad + set_value so a hand-written GD step can update params and clear grads between steps. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
This commit is contained in:
@@ -177,6 +177,41 @@ unsafe extern "C" {
|
||||
);
|
||||
}
|
||||
|
||||
// Structural ops for the tiny transformer (csrc/ops/model.cu): token embedding
|
||||
// (gather fwd / scatter-add bwd) and a 3D axis-(0,1) transpose for the multi-head
|
||||
// attention layout. F32 values, I32 ids, row-major contiguous.
|
||||
#[cfg(not(no_cuda))]
|
||||
unsafe extern "C" {
|
||||
// Embedding: out[s,:] = table[ids[s], :]. table:[vocab,dim], ids:[seq] (I32).
|
||||
pub fn launch_embedding_fwd_f32(
|
||||
table: *const f32,
|
||||
ids: *const i32,
|
||||
out: *mut f32,
|
||||
seq: i32,
|
||||
dim: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
// Scatter-add: dtable[ids[s],:] += dout[s,:] (dtable pre-zeroed; atomic).
|
||||
pub fn launch_embedding_bwd_f32(
|
||||
dout: *const f32,
|
||||
ids: *const i32,
|
||||
dtable: *mut f32,
|
||||
seq: i32,
|
||||
dim: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
|
||||
// 3D axis-(0,1) transpose: in:[a,b,c] -> out:[b,a,c]. out[j,i,k]=in[i,j,k].
|
||||
pub fn launch_transpose_3d01_f32(
|
||||
input: *const f32,
|
||||
out: *mut f32,
|
||||
a: i32,
|
||||
b: i32,
|
||||
c: i32,
|
||||
s: CudaStream,
|
||||
);
|
||||
}
|
||||
|
||||
// cuBLAS — used ONLY as a correctness reference for the hand-written GEMM in
|
||||
// tests. Declared (and linked, see build.rs) only when CUDA is compiled in.
|
||||
#[cfg(not(no_cuda))]
|
||||
|
||||
Reference in New Issue
Block a user