| use bytemuck::{Pod, Zeroable}; |
| use half::f16; |
| use std::borrow::Cow; |
| use wgpu::util::DeviceExt; |
|
|
| #[repr(C)] |
| #[derive(Copy, Clone, Pod, Zeroable)] |
| struct Params { m: u32, n: u32, k: u32, pad: u32 } |
|
|
| const WGSL: &str = r#" |
| enable f16; |
| struct Params { m: u32, n: u32, k: u32, pad: u32 }; |
| @group(0) @binding(0) var<storage, read> a: array<f16>; |
| @group(0) @binding(1) var<storage, read> b: array<f16>; |
| @group(0) @binding(2) var<storage, read_write> c: array<f16>; |
| @group(0) @binding(3) var<uniform> p: Params; |
| var<workgroup> As: array<f16, 256>; |
| var<workgroup> Bs: array<f16, 256>; |
| @compute @workgroup_size(16, 16) |
| fn main(@builtin(workgroup_id) wid: vec3<u32>, @builtin(local_invocation_id) lid: vec3<u32>) { |
| let tx = lid.x; let ty = lid.y; |
| let row0 = wid.x * 16u; let col0 = wid.y * 16u; |
| let row = row0 + ty; let col = col0 + tx; |
| var acc = 0.0; |
| var k0 = 0u; |
| loop { |
| if (k0 >= p.k) { break; } |
| if (row < p.m && (k0 + tx) < p.k) { As[ty * 16u + tx] = a[row * p.k + k0 + tx]; } else { As[ty * 16u + tx] = f16(0.0); } |
| if ((col0 + ty) < p.n && (k0 + tx) < p.k) { Bs[ty * 16u + tx] = b[(col0 + ty) * p.k + k0 + tx]; } else { Bs[ty * 16u + tx] = f16(0.0); } |
| workgroupBarrier(); |
| var kk = 0u; |
| loop { if (kk >= 16u) { break; } acc = acc + f32(As[ty * 16u + kk]) * f32(Bs[tx * 16u + kk]); kk = kk + 1u; } |
| workgroupBarrier(); |
| k0 = k0 + 16u; |
| } |
| if (row < p.m && col < p.n) { c[row * p.n + col] = f16(acc); } |
| } |
| "#; |
|
|
| fn gelu(x: f32) -> f32 { |
| 0.5 * x * (1.0 + (0.79788456 * (x + 0.044715 * x * x * x)).tanh()) |
| } |
| fn q(v: f32) -> f32 { f16::from_f32(v).to_f32() } |
|
|
| fn cpu_ffn(x: &[f32], w1: &[f32], w2: &[f32], m: usize, d: usize, f: usize) -> Vec<f32> { |
| let mut up = vec![0.0f32; m * f]; |
| for i in 0..m { for j in 0..f { |
| let mut acc = 0.0f32; |
| for k in 0..d { acc += q(x[i * d + k]) * q(w1[j * d + k]); } |
| up[i * f + j] = q(gelu(q(acc))); |
| }} |
| let mut y = vec![0.0f32; m * d]; |
| for i in 0..m { for j in 0..d { |
| let mut acc = 0.0f32; |
| for k in 0..f { acc += q(up[i * f + k]) * q(w2[j * f + k]); } |
| y[i * d + j] = q(acc); |
| }} |
| y |
| } |
|
|
| fn main() { |
| let (m, d, f) = (1560usize, 1536usize, 8960usize); |
| let raw = std::fs::read("D:/3DWord/itseez3d/ai-lab/lingbot-world/port/wgpu-rms-layer-h2.bin").expect("h2.bin"); |
| let x16: Vec<f16> = bytemuck::cast_slice::<u8, f16>(&raw)[..m * d].to_vec(); |
| let x: Vec<f32> = x16.iter().map(|v| v.to_f32()).collect(); |
| let mut w1 = vec![0.0f32; f * d]; |
| let mut w2 = vec![0.0f32; d * f]; |
|
|
| for i in 0..w1.len() { w1[i] = ((i * 17 % 97) as f32 - 48.0) / (d as f32 * 8.0); } |
| for i in 0..w2.len() { w2[i] = ((i * 23 % 89) as f32 - 44.0) / (f as f32 * 8.0); } |
| println!("cpu reference..."); |
| let raw2 = std::fs::read("D:/3DWord/itseez3d/ai-lab/lingbot-world/port/wgpu-qkv-rope-out.bin").expect("out.bin"); |
| let x2: Vec<f32> = bytemuck::cast_slice::<u8, f32>(&raw2)[..m * d].to_vec(); |
| let eraw = std::fs::read("D:/3DWord/itseez3d/ai-lab/lingbot-world/port/wgpu-time-e.bin").expect("time e"); |
| let efull: &[f32] = bytemuck::cast_slice(&eraw); |
| let e5: Vec<f32> = efull[5 * d..6 * d].to_vec(); |
| let mut expect = cpu_ffn(&x[..2 * d], &w1, &w2, 2, d, f); |
| for i in 0..2 { for j in 0..d { expect[i * d + j] = expect[i * d + j] * e5[j] + x2[i * d + j]; } } |
|
|
| let instance = wgpu::Instance::new(&wgpu::InstanceDescriptor { backends: wgpu::Backends::VULKAN, ..Default::default() }); |
| let adapter = pollster::block_on(instance.request_adapter(&wgpu::RequestAdapterOptions { |
| power_preference: wgpu::PowerPreference::HighPerformance, compatible_surface: None, force_fallback_adapter: false, |
| })).expect("no Vulkan adapter"); |
| let info = adapter.get_info(); |
| println!("adapter: {} | backend: {:?}", info.name, info.backend); |
| let mut limits = wgpu::Limits::downlevel_defaults(); |
| limits.max_storage_buffer_binding_size = 256 << 20; |
| let (device, queue) = pollster::block_on(adapter.request_device(&wgpu::DeviceDescriptor { |
| label: Some("ffn"), required_features: wgpu::Features::SHADER_F16, |
| required_limits: limits, ..Default::default() |
| })).expect("device"); |
| let mk = |data: &[f16]| device.create_buffer_init(&wgpu::util::BufferInitDescriptor { |
| label: None, contents: bytemuck::cast_slice(data), usage: wgpu::BufferUsages::STORAGE, |
| }); |
| let x16: Vec<f16> = x.iter().map(|v| f16::from_f32(*v)).collect(); |
| let w1_16: Vec<f16> = w1.iter().map(|v| f16::from_f32(*v)).collect(); |
| let w2_16: Vec<f16> = w2.iter().map(|v| f16::from_f32(*v)).collect(); |
| let x_buf = mk(&x16); let w1_buf = mk(&w1_16); let w2_buf = mk(&w2_16); |
| let ffn_buf = device.create_buffer(&wgpu::BufferDescriptor { label: None, size: (m * f * 2) as u64, usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC | wgpu::BufferUsages::COPY_DST, mapped_at_creation: false }); |
| let act_buf = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { |
| label: None, contents: bytemuck::cast_slice(&vec![f16::from_f32(0.0); m * f]), |
| usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_DST, |
| }); |
| let y_buf = device.create_buffer(&wgpu::BufferDescriptor { label: None, size: (m * d * 2) as u64, usage: wgpu::BufferUsages::STORAGE | wgpu::BufferUsages::COPY_SRC, mapped_at_creation: false }); |
| let up_p = Params { m: m as u32, n: f as u32, k: d as u32, pad: 0 }; |
| let dn_p = Params { m: m as u32, n: d as u32, k: f as u32, pad: 0 }; |
| let up_pb = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: None, contents: bytemuck::bytes_of(&up_p), usage: wgpu::BufferUsages::UNIFORM }); |
| let dn_pb = device.create_buffer_init(&wgpu::util::BufferInitDescriptor { label: None, contents: bytemuck::bytes_of(&dn_p), usage: wgpu::BufferUsages::UNIFORM }); |
| let up_rb = device.create_buffer(&wgpu::BufferDescriptor { label: None, size: (m * f * 2) as u64, usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false }); |
| let y_rb = device.create_buffer(&wgpu::BufferDescriptor { label: None, size: (m * d * 2) as u64, usage: wgpu::BufferUsages::COPY_DST | wgpu::BufferUsages::MAP_READ, mapped_at_creation: false }); |
| let module = device.create_shader_module(wgpu::ShaderModuleDescriptor { label: None, source: wgpu::ShaderSource::Wgsl(Cow::Borrowed(WGSL)) }); |
| let pipe = device.create_compute_pipeline(&wgpu::ComputePipelineDescriptor { label: None, layout: None, module: &module, entry_point: Some("main"), compilation_options: Default::default(), cache: None }); |
| let bind = |x: &wgpu::Buffer, w: &wgpu::Buffer, y: &wgpu::Buffer, pb: &wgpu::Buffer| device.create_bind_group(&wgpu::BindGroupDescriptor { label: None, layout: &pipe.get_bind_group_layout(0), entries: &[ |
| wgpu::BindGroupEntry { binding: 0, resource: x.as_entire_binding() }, |
| wgpu::BindGroupEntry { binding: 1, resource: w.as_entire_binding() }, |
| wgpu::BindGroupEntry { binding: 2, resource: y.as_entire_binding() }, |
| wgpu::BindGroupEntry { binding: 3, resource: pb.as_entire_binding() }, |
| ]}); |
| let up_bg = bind(&x_buf, &w1_buf, &ffn_buf, &up_pb); |
| let dn_bg = bind(&act_buf, &w2_buf, &y_buf, &dn_pb); |
| let mut enc = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); |
| { |
| let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None, timestamp_writes: None }); |
| p.set_pipeline(&pipe); p.set_bind_group(0, &up_bg, &[]); p.dispatch_workgroups(((m + 15) / 16) as u32, ((f + 15) / 16) as u32, 1); |
| } |
| enc.copy_buffer_to_buffer(&ffn_buf, 0, &up_rb, 0, (m * f * 2) as u64); |
| queue.submit(Some(enc.finish())); |
| let su = up_rb.slice(..); su.map_async(wgpu::MapMode::Read, |r| r.expect("map")); |
| device.poll(wgpu::PollType::wait_indefinitely()).expect("poll"); |
| let du = su.get_mapped_range(); |
| let up: &[f16] = bytemuck::cast_slice(&du); |
| let mut act = vec![f16::from_f32(0.0); m * f]; |
| for i in 0..m * f { act[i] = f16::from_f32(gelu(up[i].to_f32())); } |
| drop(du); |
| queue.write_buffer(&act_buf, 0, bytemuck::cast_slice(&act)); |
| let mut enc = device.create_command_encoder(&wgpu::CommandEncoderDescriptor { label: None }); |
| { |
| let mut p = enc.begin_compute_pass(&wgpu::ComputePassDescriptor { label: None, timestamp_writes: None }); |
| p.set_pipeline(&pipe); p.set_bind_group(0, &dn_bg, &[]); p.dispatch_workgroups(((m + 15) / 16) as u32, ((d + 15) / 16) as u32, 1); |
| } |
| enc.copy_buffer_to_buffer(&y_buf, 0, &y_rb, 0, (m * d * 2) as u64); |
| let t0 = std::time::Instant::now(); |
| queue.submit(Some(enc.finish())); |
| let sy = y_rb.slice(..); sy.map_async(wgpu::MapMode::Read, |r| r.expect("map")); |
| device.poll(wgpu::PollType::wait_indefinitely()).expect("poll"); |
| let dy = sy.get_mapped_range(); |
| let y: &[f16] = bytemuck::cast_slice(&dy); |
| let mut max_err = 0.0f32; |
| for i in 0..2 * d { |
| let got = y[i].to_f32() * e5[i % d] + x2[i]; |
| max_err = max_err.max((got - expect[i]).abs()); |
| } |
| let mut y16 = vec![f16::from_f32(0.0); m * d]; |
| for i in 0..m * d { y16[i] = f16::from_f32(y[i].to_f32() * e5[i % d] + x2[i]); } |
| std::fs::write("D:/3DWord/itseez3d/ai-lab/lingbot-world/port/wgpu-layer-y.bin", bytemuck::cast_slice(&y16)).expect("write y"); |
| println!("ffn split M={m} D={d} F={f} down_ms={:.3} max_abs_err={max_err:.3e}", t0.elapsed().as_secs_f64() * 1000.0); |
| assert!(max_err < 5e-2, "ffn mismatch"); |
| println!("VULKAN SPLIT FFN OK"); |
| }
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
