diff --git a/src/vulkan/mod.rs b/src/vulkan/mod.rs index 73da16e..ac5786f 100644 --- a/src/vulkan/mod.rs +++ b/src/vulkan/mod.rs @@ -147,6 +147,8 @@ struct Inner { gfx11: bool, /// The device runs fixed 64-lane subgroups, split into logical 32-lane warps. wave64: bool, + /// Use double-precision quantizer reciprocals (`RCP_D`, spec constant 13) on AMD devices even in wave32. + rcp_d: bool, /// VK_KHR_cooperative_matrix with the 16 x 16 x 16 f16 / f32 shape is enabled (`gemm_f`, `attn_prefill`). coop: bool, /// Without cooperative matrices: run the tiled `attn_prefill_tile` (one KV-head tile shared by 32 query/head tasks) @@ -467,7 +469,8 @@ impl Vk { prof: if std::env::var_os("STRATA_VK_PROF").is_some() { Some((unsafe { dev.create_query_pool(&vk::QueryPoolCreateInfo::default().query_type(vk::QueryType::TIMESTAMP).query_count(65536), None) }?, vec![], HashMap::new())) } else { None }, step_pos: None }; let cr = cond.then(|| ash::ext::conditional_rendering::Device::new(&instance, &dev)); let ts_period = limits.timestamp_period; - let inner = Inner { _entry: entry, instance, pd, dev, queue, tq, qfs: [qf, tf.unwrap_or(qf)], csem, tsem, mem, budget, igpu, gfx11, wave64, coop, dp_tile, coop_i8, amd_coherent, amd_plain, cr, llama: crate::cuda::llama_numerics() as u32, kv_q8: crate::cuda::kv_q8(), smem_limit, ts_period, stats, layout, st: Mutex::new(st) }; + let rcp_d = dev_props.vendor_id == 0x1002; + let inner = Inner { _entry: entry, instance, pd, dev, queue, tq, qfs: [qf, tf.unwrap_or(qf)], csem, tsem, mem, budget, igpu, gfx11, wave64, rcp_d, coop, dp_tile, coop_i8, amd_coherent, amd_plain, cr, llama: crate::cuda::llama_numerics() as u32, kv_q8: crate::cuda::kv_q8(), smem_limit, ts_period, stats, layout, st: Mutex::new(st) }; let staging = inner.buffer(STAGING, HOST)?; let map = unsafe { inner.dev.map_memory(staging.1, 0, vk::WHOLE_SIZE, vk::MemoryMapFlags::empty()) }; let mut st = inner.st(); @@ -773,8 +776,9 @@ impl Inner { let n = spec.len() as u32; let entries: Vec<_> = (0..n).map(|i| vk::SpecializationMapEntry::default().constant_id(i).offset(4 * i).size(4)) .chain([vk::SpecializationMapEntry::default().constant_id(15).offset(4 * n).size(4), - vk::SpecializationMapEntry::default().constant_id(14).offset(4 * n + 4).size(4)]).collect(); - let data: Vec = spec.iter().chain([&(self.kv_q8 as u32), &(self.wave64 as u32)]).flat_map(|v| v.to_le_bytes()).collect(); + vk::SpecializationMapEntry::default().constant_id(14).offset(4 * n + 4).size(4), + vk::SpecializationMapEntry::default().constant_id(13).offset(4 * n + 8).size(4)]).collect(); + let data: Vec = spec.iter().chain([&(self.kv_q8 as u32), &(self.wave64 as u32), &(self.rcp_d as u32)]).flat_map(|v| v.to_le_bytes()).collect(); let si = vk::SpecializationInfo::default().map_entries(&entries).data(&data); let mut sub = vk::PipelineShaderStageRequiredSubgroupSizeCreateInfo::default().required_subgroup_size(32); let stage = vk::PipelineShaderStageCreateInfo::default().stage(vk::ShaderStageFlags::COMPUTE).module(m).name(c"main").specialization_info(&si); diff --git a/src/vulkan/shaders/common.glsl b/src/vulkan/shaders/common.glsl index 6ee5af7..4fe6306 100644 --- a/src/vulkan/shaders/common.glsl +++ b/src/vulkan/shaders/common.glsl @@ -8,6 +8,8 @@ // 1 on devices whose physical subgroup is 64 lanes, split into two logical 32-lane warps. Set on every pipeline. layout(constant_id = 14) const bool W64 = false; +// 1 on AMD devices (including wave32 RADV) so quantizer reciprocals use exact double division. +layout(constant_id = 13) const bool RCP_D = false; #define LANE (W64 ? (gl_SubgroupInvocationID & 31u) : gl_SubgroupInvocationID) #define WARP (W64 ? (gl_SubgroupID * (gl_SubgroupSize / 32u) + gl_SubgroupInvocationID / 32u) : gl_SubgroupID) #define NWARPS ((gl_WorkGroupSize.x + 31u) / 32u) @@ -29,9 +31,9 @@ float fdiv(float a, float b) { // An infinite operand or b = 0 makes the steps inf * 0 = NaN (sigmoid(-100): 1 / inf); a / b is exact there. return isinf(a) || isinf(b) || b == 0.0 ? a / b : r; } -// Quantizer reciprocals on wave64: a correctly rounded double division keeps values on a roundEven() tie -// from moving by one int8 step. Wave32 keeps the fdiv path that its kernels were tuned with. -float rcp_d(float a, float b) { return W64 ? float(double(a) / double(b)) : fdiv(a, b); } +// Quantizer reciprocals on wave64 and AMD wave32: a correctly rounded double division keeps values on a +// roundEven() tie from moving by one int8 step. Non-AMD wave32 keeps the fdiv path that its kernels were tuned with. +float rcp_d(float a, float b) { return (W64 || RCP_D) ? float(double(a) / double(b)) : fdiv(a, b); } float sqrt_(float x) { float s = sqrt(x); precise float r = fma(fma(-s, s, x), fdiv(0.5, s), s); return x > 0.0 && !isinf(x) ? r : s; } // expf as CUDA compiles it (its PTX): x = j ln2 + r with j from a rounded-down fma, then ex2.approx(r) * 2^j. diff --git a/src/vulkan/spv/argmax_rows.spv b/src/vulkan/spv/argmax_rows.spv index 43bd9c2..e17cd6d 100644 Binary files a/src/vulkan/spv/argmax_rows.spv and b/src/vulkan/spv/argmax_rows.spv differ diff --git a/src/vulkan/spv/attn_decode.spv b/src/vulkan/spv/attn_decode.spv index 14cb0c9..639bddb 100644 Binary files a/src/vulkan/spv/attn_decode.spv and b/src/vulkan/spv/attn_decode.spv differ diff --git a/src/vulkan/spv/attn_decode_llama.spv b/src/vulkan/spv/attn_decode_llama.spv index ac5d29c..3020de6 100644 Binary files a/src/vulkan/spv/attn_decode_llama.spv and b/src/vulkan/spv/attn_decode_llama.spv differ diff --git a/src/vulkan/spv/attn_merge.spv b/src/vulkan/spv/attn_merge.spv index 65c829d..bd78efc 100644 Binary files a/src/vulkan/spv/attn_merge.spv and b/src/vulkan/spv/attn_merge.spv differ diff --git a/src/vulkan/spv/attn_prefill_128.spv b/src/vulkan/spv/attn_prefill_128.spv index bab8012..b03c4ec 100644 Binary files a/src/vulkan/spv/attn_prefill_128.spv and b/src/vulkan/spv/attn_prefill_128.spv differ diff --git a/src/vulkan/spv/attn_prefill_256.spv b/src/vulkan/spv/attn_prefill_256.spv index 18386b0..dde15c7 100644 Binary files a/src/vulkan/spv/attn_prefill_256.spv and b/src/vulkan/spv/attn_prefill_256.spv differ diff --git a/src/vulkan/spv/attn_prefill_64.spv b/src/vulkan/spv/attn_prefill_64.spv index 904c280..ebb2715 100644 Binary files a/src/vulkan/spv/attn_prefill_64.spv and b/src/vulkan/spv/attn_prefill_64.spv differ diff --git a/src/vulkan/spv/attn_prefill_dp.spv b/src/vulkan/spv/attn_prefill_dp.spv index 2f3d50a..b1c7e2d 100644 Binary files a/src/vulkan/spv/attn_prefill_dp.spv and b/src/vulkan/spv/attn_prefill_dp.spv differ diff --git a/src/vulkan/spv/attn_prefill_llama.spv b/src/vulkan/spv/attn_prefill_llama.spv index 819350b..a1fb9b6 100644 Binary files a/src/vulkan/spv/attn_prefill_llama.spv and b/src/vulkan/spv/attn_prefill_llama.spv differ diff --git a/src/vulkan/spv/attn_prefill_tile.spv b/src/vulkan/spv/attn_prefill_tile.spv index a4cb598..2af1400 100644 Binary files a/src/vulkan/spv/attn_prefill_tile.spv and b/src/vulkan/spv/attn_prefill_tile.spv differ diff --git a/src/vulkan/spv/conv_seq.spv b/src/vulkan/spv/conv_seq.spv index c76a2c2..692f646 100644 Binary files a/src/vulkan/spv/conv_seq.spv and b/src/vulkan/spv/conv_seq.spv differ diff --git a/src/vulkan/spv/conv_silu.spv b/src/vulkan/spv/conv_silu.spv index 607904f..0976bf9 100644 Binary files a/src/vulkan/spv/conv_silu.spv and b/src/vulkan/spv/conv_silu.spv differ diff --git a/src/vulkan/spv/conv_state.spv b/src/vulkan/spv/conv_state.spv index ad2a6af..75ea35f 100644 Binary files a/src/vulkan/spv/conv_state.spv and b/src/vulkan/spv/conv_state.spv differ diff --git a/src/vulkan/spv/deltanet.spv b/src/vulkan/spv/deltanet.spv index 45b8a3a..a35be9d 100644 Binary files a/src/vulkan/spv/deltanet.spv and b/src/vulkan/spv/deltanet.spv differ diff --git a/src/vulkan/spv/dequant.spv b/src/vulkan/spv/dequant.spv index a207e30..d3c5e12 100644 Binary files a/src/vulkan/spv/dequant.spv and b/src/vulkan/spv/dequant.spv differ diff --git a/src/vulkan/spv/gemm_f.spv b/src/vulkan/spv/gemm_f.spv index 981e103..3e29c5d 100644 Binary files a/src/vulkan/spv/gemm_f.spv and b/src/vulkan/spv/gemm_f.spv differ diff --git a/src/vulkan/spv/gemm_f_dp.spv b/src/vulkan/spv/gemm_f_dp.spv index 7a03d76..6d35489 100644 Binary files a/src/vulkan/spv/gemm_f_dp.spv and b/src/vulkan/spv/gemm_f_dp.spv differ diff --git a/src/vulkan/spv/gemm_fs.spv b/src/vulkan/spv/gemm_fs.spv index dddf75d..af7b380 100644 Binary files a/src/vulkan/spv/gemm_fs.spv and b/src/vulkan/spv/gemm_fs.spv differ diff --git a/src/vulkan/spv/gemmi.spv b/src/vulkan/spv/gemmi.spv index aa102a5..0df03ad 100644 Binary files a/src/vulkan/spv/gemmi.spv and b/src/vulkan/spv/gemmi.spv differ diff --git a/src/vulkan/spv/gemmi_dp.spv b/src/vulkan/spv/gemmi_dp.spv index 962b6e7..a2fbd1b 100644 Binary files a/src/vulkan/spv/gemmi_dp.spv and b/src/vulkan/spv/gemmi_dp.spv differ diff --git a/src/vulkan/spv/gemmi_s.spv b/src/vulkan/spv/gemmi_s.spv index 8df0d35..2e22696 100644 Binary files a/src/vulkan/spv/gemmi_s.spv and b/src/vulkan/spv/gemmi_s.spv differ diff --git a/src/vulkan/spv/gemv_f.spv b/src/vulkan/spv/gemv_f.spv index 3a11800..aa3646b 100644 Binary files a/src/vulkan/spv/gemv_f.spv and b/src/vulkan/spv/gemv_f.spv differ diff --git a/src/vulkan/spv/gemv_q.spv b/src/vulkan/spv/gemv_q.spv index 9cf8eb2..5348b32 100644 Binary files a/src/vulkan/spv/gemv_q.spv and b/src/vulkan/spv/gemv_q.spv differ diff --git a/src/vulkan/spv/hc_mix.spv b/src/vulkan/spv/hc_mix.spv index 0bb73c1..e213f7e 100644 Binary files a/src/vulkan/spv/hc_mix.spv and b/src/vulkan/spv/hc_mix.spv differ diff --git a/src/vulkan/spv/hc_silu.spv b/src/vulkan/spv/hc_silu.spv index 529937a..437cbd7 100644 Binary files a/src/vulkan/spv/hc_silu.spv and b/src/vulkan/spv/hc_silu.spv differ diff --git a/src/vulkan/spv/hc_write.spv b/src/vulkan/spv/hc_write.spv index ea80053..e529e0a 100644 Binary files a/src/vulkan/spv/hc_write.spv and b/src/vulkan/spv/hc_write.spv differ diff --git a/src/vulkan/spv/kv_store.spv b/src/vulkan/spv/kv_store.spv index ae1fa79..4a892f8 100644 Binary files a/src/vulkan/spv/kv_store.spv and b/src/vulkan/spv/kv_store.spv differ diff --git a/src/vulkan/spv/l2_norm.spv b/src/vulkan/spv/l2_norm.spv index 770eb07..d24fb72 100644 Binary files a/src/vulkan/spv/l2_norm.spv and b/src/vulkan/spv/l2_norm.spv differ diff --git a/src/vulkan/spv/moe_combine.spv b/src/vulkan/spv/moe_combine.spv index c8f72c7..a88b03a 100644 Binary files a/src/vulkan/spv/moe_combine.spv and b/src/vulkan/spv/moe_combine.spv differ diff --git a/src/vulkan/spv/moe_combine_seq.spv b/src/vulkan/spv/moe_combine_seq.spv index 579136e..5049671 100644 Binary files a/src/vulkan/spv/moe_combine_seq.spv and b/src/vulkan/spv/moe_combine_seq.spv differ diff --git a/src/vulkan/spv/moe_down_combine.spv b/src/vulkan/spv/moe_down_combine.spv index 8069f06..f79c755 100644 Binary files a/src/vulkan/spv/moe_down_combine.spv and b/src/vulkan/spv/moe_down_combine.spv differ diff --git a/src/vulkan/spv/moe_gemv_f.spv b/src/vulkan/spv/moe_gemv_f.spv index d341d82..e0d6041 100644 Binary files a/src/vulkan/spv/moe_gemv_f.spv and b/src/vulkan/spv/moe_gemv_f.spv differ diff --git a/src/vulkan/spv/moe_gemv_q.spv b/src/vulkan/spv/moe_gemv_q.spv index c77b1bb..6ffd4de 100644 Binary files a/src/vulkan/spv/moe_gemv_q.spv and b/src/vulkan/spv/moe_gemv_q.spv differ diff --git a/src/vulkan/spv/moe_route.spv b/src/vulkan/spv/moe_route.spv index b741d87..904ebc7 100644 Binary files a/src/vulkan/spv/moe_route.spv and b/src/vulkan/spv/moe_route.spv differ diff --git a/src/vulkan/spv/norm_gate.spv b/src/vulkan/spv/norm_gate.spv index 69aa101..bdce630 100644 Binary files a/src/vulkan/spv/norm_gate.spv and b/src/vulkan/spv/norm_gate.spv differ diff --git a/src/vulkan/spv/ple_conv.spv b/src/vulkan/spv/ple_conv.spv index af1ffac..b80ce58 100644 Binary files a/src/vulkan/spv/ple_conv.spv and b/src/vulkan/spv/ple_conv.spv differ diff --git a/src/vulkan/spv/ple_gate.spv b/src/vulkan/spv/ple_gate.spv index 5fcdaac..c7cb065 100644 Binary files a/src/vulkan/spv/ple_gate.spv and b/src/vulkan/spv/ple_gate.spv differ diff --git a/src/vulkan/spv/qsa_pool.spv b/src/vulkan/spv/qsa_pool.spv index 86cff62..51da4cb 100644 Binary files a/src/vulkan/spv/qsa_pool.spv and b/src/vulkan/spv/qsa_pool.spv differ diff --git a/src/vulkan/spv/qsa_score.spv b/src/vulkan/spv/qsa_score.spv index 1311d16..feb5128 100644 Binary files a/src/vulkan/spv/qsa_score.spv and b/src/vulkan/spv/qsa_score.spv differ diff --git a/src/vulkan/spv/qsa_select.spv b/src/vulkan/spv/qsa_select.spv index 88b07cb..7b04774 100644 Binary files a/src/vulkan/spv/qsa_select.spv and b/src/vulkan/spv/qsa_select.spv differ diff --git a/src/vulkan/spv/quant_q8.spv b/src/vulkan/spv/quant_q8.spv index e673917..76cc59b 100644 Binary files a/src/vulkan/spv/quant_q8.spv and b/src/vulkan/spv/quant_q8.spv differ diff --git a/src/vulkan/spv/quant_xq.spv b/src/vulkan/spv/quant_xq.spv index 746c834..e9f9afe 100644 Binary files a/src/vulkan/spv/quant_xq.spv and b/src/vulkan/spv/quant_xq.spv differ diff --git a/src/vulkan/spv/rms_norm.spv b/src/vulkan/spv/rms_norm.spv index 2fa6676..cb6e1fb 100644 Binary files a/src/vulkan/spv/rms_norm.spv and b/src/vulkan/spv/rms_norm.spv differ diff --git a/src/vulkan/spv/rope_neox.spv b/src/vulkan/spv/rope_neox.spv index 01ad66f..9e7a2ee 100644 Binary files a/src/vulkan/spv/rope_neox.spv and b/src/vulkan/spv/rope_neox.spv differ diff --git a/src/vulkan/spv/router_topk.spv b/src/vulkan/spv/router_topk.spv index 85813ca..a000dbe 100644 Binary files a/src/vulkan/spv/router_topk.spv and b/src/vulkan/spv/router_topk.spv differ diff --git a/src/vulkan/spv/ssm_gates.spv b/src/vulkan/spv/ssm_gates.spv index e3a9095..4633a73 100644 Binary files a/src/vulkan/spv/ssm_gates.spv and b/src/vulkan/spv/ssm_gates.spv differ diff --git a/src/vulkan/spv/swiglu.spv b/src/vulkan/spv/swiglu.spv index 8cfa923..20b5de4 100644 Binary files a/src/vulkan/spv/swiglu.spv and b/src/vulkan/spv/swiglu.spv differ