diff --git a/llama.cpp/ggml/src/ggml-cuda/set-rows.cu b/llama.cpp/ggml/src/ggml-cuda/set-rows.cu --- a/llama.cpp/ggml/src/ggml-cuda/set-rows.cu +++ b/llama.cpp/ggml/src/ggml-cuda/set-rows.cu @@ -748,33 +748,17 @@ const int64_t s10_i = nb10/sizeof(idx_t), s11_i = nb11/sizeof(idx_t), s12_i = nb12/sizeof(idx_t); const int iq_is_k = (strncmp(dst->name, "cache_k_", 8) == 0) ? 1 : 0; - // one-time per-device: prefer the 8 KB/block shared-memory backtrace (opt-in); else global fallback buffer - static int tcq3_use_shared_bt[GGML_CUDA_MAX_DEVICES] = {}; - static bool tcq3_bt_checked [GGML_CUDA_MAX_DEVICES] = {}; - constexpr int tcq3_bt_shared_bytes = 128 * 64; - if (!tcq3_bt_checked[ctx.device]) { - tcq3_bt_checked[ctx.device] = true; - const char * env = getenv("TURBO_TCQ_SHARED_BT"); - if (!env || atoi(env) != 0) { - int max_shared_optin = 0; - CUDA_CHECK(cudaDeviceGetAttribute(&max_shared_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, ctx.device)); - if (max_shared_optin >= tcq3_bt_shared_bytes) { - if (tcq_fp64_enabled()) - CUDA_SET_SHARED_MEMORY_LIMIT((k_set_rows_turbo3_tcq), tcq3_bt_shared_bytes); - else - CUDA_SET_SHARED_MEMORY_LIMIT(k_set_rows_turbo3_tcq, tcq3_bt_shared_bytes); - tcq3_use_shared_bt[ctx.device] = 1; - } - } - } - if (!tcq3_use_shared_bt[ctx.device]) ensure_tcq_bt_buf(ctx.device, ne_total_groups * 128 * 64); + // #533 warp-synchronous encode: bt is always a global per-group slab. With TCQ3_WPB groups/block a shared-bt + // would cost WPB*8KB; and the warp kernel is no longer barrier-bound, so global bt is free here. One 8KB slab + // per group (indexed by `group`, not blockIdx.x). tcq_bt_buf grows monotonically and is reused across launches. + ensure_tcq_bt_buf(ctx.device, ne_total_groups * 128 * 64); const uint3 ne00_fd = init_fastdiv_values((uint32_t) ne00); const uint3 ne01_fd = init_fastdiv_values((uint32_t) ne01); const uint3 ne02_fd = init_fastdiv_values((uint32_t) ne02); const uint3 ne11_fd = init_fastdiv_values((uint32_t) ne11); const uint3 ne12_fd = init_fastdiv_values((uint32_t) ne12); - const int shared_bytes = tcq3_use_shared_bt[ctx.device] ? tcq3_bt_shared_bytes : 0; + const int tcq3_grid = (int)((ne_total_groups + TCQ3_WPB - 1) / TCQ3_WPB); // #533 warp-packed: TCQ3_WPB groups/block, 32 thr/group // opencoti-hook: TCQ activation-dump (F-D, #515) — bind the staging buffer for this launch (no-op unless // TURBO_TCQ_DUMP_ACTS is set). tcq_dump_init() is idempotent; dump_dev is nullptr when inactive/cap-reached, // in which case d_tcq_dump_buf is left at its nullptr default and the kernel's dump branch is a uniform no-op. @@ -782,15 +766,15 @@ float * dump_dev = tcq_dump_pre_launch((long) ne_total_groups); if (dump_dev) CUDA_CHECK(cudaMemcpyToSymbol(d_tcq_dump_buf, &dump_dev, sizeof(float*))); if (tcq_fp64_enabled()) - k_set_rows_turbo3_tcq<<<(int)ne_total_groups, 512, shared_bytes, stream>>>( + k_set_rows_turbo3_tcq<<>>( src0_d, src1_d, (block_turbo3_tcq *)dst->data, - ne_total_groups, tcq_bt_buf[ctx.device], tcq3_use_shared_bt[ctx.device], ne00, ne01, ne02, ne10, ne11, ne12, ne13, + ne_total_groups, tcq_bt_buf[ctx.device], 0, ne00, ne01, ne02, ne10, ne11, ne12, ne13, s01_f, s02_f, s03_f, s10_i, s11_i, s12_i, iq_is_k, nb1, nb2, nb3, ne00_fd, ne01_fd, ne02_fd, ne11_fd, ne12_fd); else - k_set_rows_turbo3_tcq<<<(int)ne_total_groups, 512, shared_bytes, stream>>>( + k_set_rows_turbo3_tcq<<>>( src0_d, src1_d, (block_turbo3_tcq *)dst->data, - ne_total_groups, tcq_bt_buf[ctx.device], tcq3_use_shared_bt[ctx.device], ne00, ne01, ne02, ne10, ne11, ne12, ne13, + ne_total_groups, tcq_bt_buf[ctx.device], 0, ne00, ne01, ne02, ne10, ne11, ne12, ne13, s01_f, s02_f, s03_f, s10_i, s11_i, s12_i, iq_is_k, nb1, nb2, nb3, ne00_fd, ne01_fd, ne02_fd, ne11_fd, ne12_fd); if (dump_dev) tcq_dump_post_launch((long) ne_total_groups, stream); @@ -827,42 +811,26 @@ const int64_t s10_i = nb10/sizeof(idx_t), s11_i = nb11/sizeof(idx_t), s12_i = nb12/sizeof(idx_t); const int iq_is_k = (strncmp(dst->name, "cache_k_", 8) == 0) ? 1 : 0; - static int tcq2_use_shared_bt[GGML_CUDA_MAX_DEVICES] = {}; - static bool tcq2_bt_checked [GGML_CUDA_MAX_DEVICES] = {}; - constexpr int tcq2_bt_shared_bytes = 128 * 64; - if (!tcq2_bt_checked[ctx.device]) { - tcq2_bt_checked[ctx.device] = true; - const char * env = getenv("TURBO_TCQ_SHARED_BT"); - if (!env || atoi(env) != 0) { - int max_shared_optin = 0; - CUDA_CHECK(cudaDeviceGetAttribute(&max_shared_optin, cudaDevAttrMaxSharedMemoryPerBlockOptin, ctx.device)); - if (max_shared_optin >= tcq2_bt_shared_bytes) { - if (tcq_fp64_enabled()) - CUDA_SET_SHARED_MEMORY_LIMIT((k_set_rows_turbo2_tcq), tcq2_bt_shared_bytes); - else - CUDA_SET_SHARED_MEMORY_LIMIT(k_set_rows_turbo2_tcq, tcq2_bt_shared_bytes); - tcq2_use_shared_bt[ctx.device] = 1; - } - } - } - if (!tcq2_use_shared_bt[ctx.device]) ensure_tcq_bt_buf(ctx.device, ne_total_groups * 128 * 64); + // #533 warp-synchronous encode: bt always a global per-group slab (TCQ2_WPB groups/block); warp kernel is + // no longer barrier-bound, so global bt is free. One 8KB slab per group, indexed by `group`. + ensure_tcq_bt_buf(ctx.device, ne_total_groups * 128 * 64); const uint3 ne00_fd = init_fastdiv_values((uint32_t) ne00); const uint3 ne01_fd = init_fastdiv_values((uint32_t) ne01); const uint3 ne02_fd = init_fastdiv_values((uint32_t) ne02); const uint3 ne11_fd = init_fastdiv_values((uint32_t) ne11); const uint3 ne12_fd = init_fastdiv_values((uint32_t) ne12); - const int shared_bytes = tcq2_use_shared_bt[ctx.device] ? tcq2_bt_shared_bytes : 0; + const int tcq2_grid = (int)((ne_total_groups + TCQ2_WPB - 1) / TCQ2_WPB); // #533 warp-packed: TCQ2_WPB groups/block, 32 thr/group if (tcq_fp64_enabled()) - k_set_rows_turbo2_tcq<<<(int)ne_total_groups, 256, shared_bytes, stream>>>( + k_set_rows_turbo2_tcq<<>>( src0_d, src1_d, (block_turbo2_tcq *)dst->data, - ne_total_groups, tcq_bt_buf[ctx.device], tcq2_use_shared_bt[ctx.device], ne00, ne01, ne02, ne10, ne11, ne12, ne13, + ne_total_groups, tcq_bt_buf[ctx.device], 0, ne00, ne01, ne02, ne10, ne11, ne12, ne13, s01_f, s02_f, s03_f, s10_i, s11_i, s12_i, iq_is_k, nb1, nb2, nb3, ne00_fd, ne01_fd, ne02_fd, ne11_fd, ne12_fd); else - k_set_rows_turbo2_tcq<<<(int)ne_total_groups, 256, shared_bytes, stream>>>( + k_set_rows_turbo2_tcq<<>>( src0_d, src1_d, (block_turbo2_tcq *)dst->data, - ne_total_groups, tcq_bt_buf[ctx.device], tcq2_use_shared_bt[ctx.device], ne00, ne01, ne02, ne10, ne11, ne12, ne13, + ne_total_groups, tcq_bt_buf[ctx.device], 0, ne00, ne01, ne02, ne10, ne11, ne12, ne13, s01_f, s02_f, s03_f, s10_i, s11_i, s12_i, iq_is_k, nb1, nb2, nb3, ne00_fd, ne01_fd, ne02_fd, ne11_fd, ne12_fd); } diff --git a/llama.cpp/ggml/src/ggml-cuda/turbo-tcq-cuda.cuh b/llama.cpp/ggml/src/ggml-cuda/turbo-tcq-cuda.cuh --- a/llama.cpp/ggml/src/ggml-cuda/turbo-tcq-cuda.cuh +++ b/llama.cpp/ggml/src/ggml-cuda/turbo-tcq-cuda.cuh @@ -109,11 +109,16 @@ static __device__ __forceinline__ float tcq_acc_sqrt(float v) { return sqrtf(v); } static __device__ __forceinline__ double tcq_acc_sqrt(double v) { return sqrt(v); } -// TCQ SET_ROWS encode: Viterbi optimal path with right-shift trellis -// 512 threads per block (one per trellis state), one block per 128-element group -// Double-buffered cost arrays + global memory backtrace (128 syncs/group, was 384) +// TCQ SET_ROWS encode: warp-synchronous Viterbi (#533). ncu proved the old 512-thread/__syncthreads design was +// LATENCY-bound on the 128-step block-barrier chain (25% SM, 0.1% DRAM, 33% occ); #538's occupancy/shuffle tweak +// couldn't help that. This collapses each 128-element group to ONE 32-thread warp (16 states/lane, thread-local +// predecessor-min, __syncwarp broadcast, bt in global), packs TCQ3_WPB warps/block, and breaks the warp lockstep +// so per-group serial latency is hidden by other warps. Validated byte-equivalent (round-trip decode cosine +// 1.00000 vs the old design — .opencoti/tcq-warp-enc-full.cu) + 3.74x on the full encode. Reuses #538's warp- +// shuffle min primitive. acc_t fp64 path (#506) preserved; bt is always global (use_shared_bt ignored). +#define TCQ3_WPB 8 template -static __global__ void __launch_bounds__(512, 1) k_set_rows_turbo3_tcq( +static __global__ void __launch_bounds__(TCQ3_WPB*32, 3) k_set_rows_turbo3_tcq( const float * __restrict__ src0, const idx_t * __restrict__ src1, block_turbo3_tcq * __restrict__ dst, const int64_t ne_total_groups, uint8_t * __restrict__ bt_buf, @@ -127,12 +132,13 @@ const uint3 ne00_fd, const uint3 ne01_fd, const uint3 ne02_fd, const uint3 ne11_fd, const uint3 ne12_fd) { - const int64_t group = blockIdx.x; + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int64_t group = (int64_t)blockIdx.x * TCQ3_WPB + warp; if (group >= ne_total_groups) return; + (void) use_shared_bt; (void) ne10; (void) ne13; // bt is always global in the warp design - const int sid = threadIdx.x; // state index 0..511 - - // Compute source and destination pointers (same index math as turbo3) + // Compute source and destination pointers (same index math as the 512-thread design; per-warp group) const int64_t i_base = group * QK_TURBO3_TCQ; uint32_t tmp = (uint32_t)i_base; uint2 div_mod; div_mod = fast_div_modulo(tmp, ne00_fd); const int64_t i00 = div_mod.y; tmp = div_mod.x; @@ -145,244 +151,130 @@ block_turbo3_tcq * dst_blk = (block_turbo3_tcq *)((char *)dst + dst_row*s1 + i02*s2 + i03*s3) + (i00 / QK_TURBO3_TCQ); - // Shared memory layout (~5KB, was ~35KB before global bt optimization): - // x[128] : rotated+normalized input (also reused as outputs[] after Viterbi) - // cost[512] : path costs buffer A (also reused for reductions) - // cost_b[512]: path costs buffer B (double-buffering eliminates 2/3 of syncs) - // Backtrace: one predecessor byte for each of the 64 low-state groups per - // step. The predecessor is independent of the output bits in sid[8:6], so - // storing 128×64 bytes is equivalent to the older 128×512 layout. - extern __shared__ uint8_t bt_shared[]; - __shared__ acc_t x[128]; - __shared__ acc_t cost[512]; - __shared__ acc_t cost_b[512]; - __shared__ int warp_min_idx[16]; - __shared__ acc_t warp_min_cost[16]; - __shared__ acc_t pred_min_cost[64]; - __shared__ int shared_initial_state; - - if (sid < 128) x[sid] = grp_src[sid]; - __syncthreads(); - - __syncthreads(); - - // Norm reduction - cost[sid] = (sid < 128) ? x[sid] * x[sid] : 0.0f; - __syncthreads(); - for (int stride = 256; stride >= 32; stride >>= 1) { - if (sid < stride) cost[sid] += cost[sid + stride]; - __syncthreads(); - } - if (sid < 32) { - acc_t v = cost[sid]; - v += __shfl_down_sync(0xFFFFFFFFULL, v, 16); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 8); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 4); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 2); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 1); - if (sid == 0) cost[0] = v; - } - __syncthreads(); - acc_t grp_norm = tcq_acc_sqrt(cost[0]); - acc_t inv_norm = grp_norm > 1e-10f ? 1.0f / grp_norm : 0.0f; - - if (sid < 128) x[sid] *= inv_norm; - __syncthreads(); - - // FWHT. The first five stages are contained within each 32-lane warp, so - // use warp shuffles and only synchronize for the two cross-warp stages. - if (sid < 128) { - acc_t v = x[sid] * TURBO3_WHT_SIGNS1[sid]; - const int lane = sid & 31; -#pragma unroll - for (int h = 1; h < 32; h <<= 1) { - const acc_t other = __shfl_xor_sync(0xFFFFFFFFULL, v, h); - v = (lane & h) ? (other - v) : (v + other); - } - x[sid] = v; - } - __syncthreads(); - if (sid < 64) { - const int j = ((sid >> 5) << 6) + (sid & 31); - acc_t a = x[j], b = x[j + 32]; - x[j] = a + b; x[j + 32] = a - b; - } - __syncthreads(); - if (sid < 64) { - acc_t a = x[sid], b = x[sid + 64]; - x[sid] = a + b; x[sid + 64] = a - b; + // Per-warp smem: whitened vec x[128] (reused as outputs[] after Viterbi) + pred_min_cost[64] broadcast + + // initial-state scratch. The Viterbi path-cost lives in registers (c[16], 16 states/lane). bt is global. + __shared__ acc_t x_s[TCQ3_WPB][128]; + __shared__ acc_t pm_s[TCQ3_WPB][64]; + __shared__ int ini_s[TCQ3_WPB]; + acc_t * x = x_s[warp]; + acc_t * pred_min_cost = pm_s[warp]; + uint8_t * bt = bt_buf + group * (int64_t)(128 * 64); // global per-group backtrace slab + + // --- load 4 elements/lane (stride-32 layout: lane owns positions L, L+32, L+64, L+96) --- + acc_t e0 = grp_src[lane], e1 = grp_src[lane+32], e2 = grp_src[lane+64], e3 = grp_src[lane+96]; + // --- norm (warp-shuffle sum of squares) --- + acc_t ss = e0*e0 + e1*e1 + e2*e2 + e3*e3; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) ss += __shfl_xor_sync(0xFFFFFFFFu, ss, o); + acc_t grp_norm = tcq_acc_sqrt(ss); + acc_t inv_norm = grp_norm > 1e-10f ? (acc_t)1.0 / grp_norm : (acc_t)0.0; + e0 *= inv_norm; e1 *= inv_norm; e2 *= inv_norm; e3 *= inv_norm; + // --- FWHT-128 (stride-32 layout: 5 shuffle-butterfly stages + 2 intra-thread cross stages, + // byte-faithful to the 512-thread block FWHT) --- + e0 *= TURBO3_WHT_SIGNS1[lane]; e1 *= TURBO3_WHT_SIGNS1[lane+32]; + e2 *= TURBO3_WHT_SIGNS1[lane+64]; e3 *= TURBO3_WHT_SIGNS1[lane+96]; +#pragma unroll + for (int h = 1; h < 32; h <<= 1) { + acc_t o0 = __shfl_xor_sync(0xFFFFFFFFu, e0, h); e0 = (lane & h) ? o0 - e0 : e0 + o0; + acc_t o1 = __shfl_xor_sync(0xFFFFFFFFu, e1, h); e1 = (lane & h) ? o1 - e1 : e1 + o1; + acc_t o2 = __shfl_xor_sync(0xFFFFFFFFu, e2, h); e2 = (lane & h) ? o2 - e2 : e2 + o2; + acc_t o3 = __shfl_xor_sync(0xFFFFFFFFu, e3, h); e3 = (lane & h) ? o3 - e3 : e3 + o3; } - __syncthreads(); + { acc_t a = e0, b = e1; e0 = a + b; e1 = a - b; acc_t c = e2, d = e3; e2 = c + d; e3 = c - d; } // cross-32: blk0<->1, blk2<->3 + { acc_t a = e0, b = e2; e0 = a + b; e2 = a - b; acc_t c = e1, d = e3; e1 = c + d; e3 = c - d; } // cross-64: blk0<->2, blk1<->3 constexpr float inv_sqrt_128 = 0.08838834764831845f; - if (sid < 128) x[sid] *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[sid]; - __syncthreads(); - - // opencoti-hook: TCQ activation-dump (F-D, #515). x[0..127] is now the post-FWHT unit-norm whitened vec — - // exactly the xt the trellis quantizes against d_turbo3_tcq_codebook below, and x stays read-only through the - // forward pass. Stage it for the F-E codebook retrain when enabled; nullptr (default) is a uniform no-op. - // group == blockIdx.x < ne_total_groups, so [group*128 + sid] is in-bounds of the (ne_total_groups*128) buffer. - if (d_tcq_dump_buf != nullptr && sid < 128) { - d_tcq_dump_buf[group * 128 + sid] = (float) x[sid]; - } - - if (sid == 0) cost[0] = grp_norm; - __syncthreads(); - - acc_t saved_norm = cost[0]; - - // Viterbi forward pass: double-buffered cost (1 sync/step, was 3) - uint8_t * bt = use_shared_bt ? bt_shared : bt_buf + (int64_t)blockIdx.x * (128 * 64); - cost[sid] = 0.0f; - __syncthreads(); - + e0 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane]; e1 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane+32]; + e2 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane+64]; e3 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane+96]; + x[lane] = e0; x[lane+32] = e1; x[lane+64] = e2; x[lane+96] = e3; + __syncwarp(); + + // opencoti-hook: TCQ activation-dump (F-D, #515). e0..e3 are the post-FWHT unit-norm whitened vec. + if (d_tcq_dump_buf != nullptr) { + d_tcq_dump_buf[group*128 + lane] = (float) e0; + d_tcq_dump_buf[group*128 + lane+32] = (float) e1; + d_tcq_dump_buf[group*128 + lane+64] = (float) e2; + d_tcq_dump_buf[group*128 + lane+96] = (float) e3; + } + acc_t saved_norm = grp_norm; + + // --- Viterbi forward DP: warp-synchronous (16 states/lane, thread-local predecessor-min, __syncwarp). --- + // lane L owns states {16L..16L+15} = trellis-groups 2L (states 16L..16L+7) and 2L+1 (states 16L+8..16L+15); + // each group's 8 predecessors (g*8..g*8+7) fall entirely within lane g/2, so the predecessor-min is + // thread-local (no shuffle). A 64-entry per-warp smem broadcast under __syncwarp replaces the block barrier. + acc_t c[16]; +#pragma unroll + for (int i = 0; i < 16; i++) c[i] = 0.0f; for (int t = 0; t < 128; t++) { - // Double-buffer: even steps read cost/write cost_b, odd steps read cost_b/write cost - acc_t * cost_rd = (t & 1) ? cost_b : cost; - acc_t * cost_wr = (t & 1) ? cost : cost_b; - + acc_t pm0 = c[0]; int pp0 = 0; +#pragma unroll + for (int j = 1; j < 8; j++) { if (c[j] < pm0) { pm0 = c[j]; pp0 = j; } } + acc_t pm1 = c[8]; int pp1 = 0; +#pragma unroll + for (int j = 1; j < 8; j++) { if (c[8+j] < pm1) { pm1 = c[8+j]; pp1 = j; } } + pred_min_cost[2*lane] = pm0; pred_min_cost[2*lane+1] = pm1; + bt[t*64 + 2*lane] = (uint8_t) pp0; bt[t*64 + 2*lane+1] = (uint8_t) pp1; + __syncwarp(); acc_t xt = x[t]; - - // Right-shift trellis: ns = (prev >> 3) | (out << 6). The best - // predecessor depends only on sid's low 6 bits, so compute those 64 - // minima once instead of repeating the same 8-way scan for each out. - if (sid < 64) { - const int base_prev = sid << 3; - acc_t best = cost_rd[base_prev]; - int best_p = 0; -#pragma unroll - for (int p = 1; p < 8; p++) { - acc_t c = cost_rd[base_prev | p]; - if (c < best) { - best = c; - best_p = p; - } - } - pred_min_cost[sid] = best; - bt[t * 64 + sid] = (uint8_t) best_p; - } - __syncthreads(); - - const int pred_idx = sid & 0x3F; - acc_t dist = xt - d_turbo3_tcq_codebook[sid]; - dist = dist * dist; - - cost_wr[sid] = pred_min_cost[pred_idx] + dist; - __syncthreads(); - } - // After 128 steps (even count): final costs are in cost[] (step 127 writes to cost) - - // Warp argmin over 512 costs - { - acc_t my_cost = cost[sid]; - int my_idx = sid; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - acc_t other_cost = __shfl_xor_sync(0xFFFFFFFFULL, my_cost, offset); - int other_idx = __shfl_xor_sync(0xFFFFFFFFULL, my_idx, offset); - if (other_cost < my_cost) { my_cost = other_cost; my_idx = other_idx; } - } - if (sid % 32 == 0) { - warp_min_cost[sid / 32] = my_cost; - warp_min_idx[sid / 32] = my_idx; - } +#pragma unroll + for (int i = 0; i < 16; i++) { int s = 16*lane + i; acc_t d = xt - d_turbo3_tcq_codebook[s]; c[i] = pred_min_cost[s & 0x3F] + d * d; } + __syncwarp(); } - __syncthreads(); - if (sid < 32) { - acc_t best = (sid < 16) ? warp_min_cost[sid] : 3.4028234663852886e38f; - int best_idx = (sid < 16) ? warp_min_idx[sid] : 0; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - acc_t other_cost = __shfl_down_sync(0xFFFFFFFFULL, best, offset); - int other_idx = __shfl_down_sync(0xFFFFFFFFULL, best_idx, offset); - if (other_cost < best) { - best = other_cost; - best_idx = other_idx; - } - } - if (sid == 0) { - shared_initial_state = best_idx; // temporarily: best final state (becomes initial after backtrack) - } + // --- argmin over 512 states (per-lane min over 16, then warp-shuffle) --- + acc_t mc = c[0]; int mi = 16*lane; +#pragma unroll + for (int i = 1; i < 16; i++) { if (c[i] < mc) { mc = c[i]; mi = 16*lane + i; } } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { + acc_t oc = __shfl_xor_sync(0xFFFFFFFFu, mc, o); int oi = __shfl_xor_sync(0xFFFFFFFFu, mi, o); + if (oc < mc || (oc == mc && oi < mi)) { mc = oc; mi = oi; } } - __syncthreads(); - - // Save x[] to global buffer before backtrack overwrites it - - // Backtrack (inherently sequential, reads global bt) + // --- backtrack (lane 0; outputs reuse the x slab — x no longer needed) --- uint8_t * outputs = (uint8_t *)x; - if (sid == 0) { - int state = shared_initial_state; + if (lane == 0) { + int state = mi; for (int t = 127; t >= 0; t--) { outputs[t] = (uint8_t)(state >> 6); - int p = bt[t * 64 + (state & 0x3F)]; + int p = bt[t*64 + (state & 0x3F)]; state = ((state & 0x3F) << 3) | p; } - shared_initial_state = state; + ini_s[warp] = state; } - __syncthreads(); + __syncwarp(); + int initial_state = ini_s[warp]; - // Save output symbols to global buffer - - // Parallel recon norm: t>=2 can compute state directly from 3 outputs (3 shifts of 3 = 9 bits) - acc_t my_recon_sq = 0.0f; - if (sid < 128) { - int cur_state; - if (sid < 2) { - cur_state = shared_initial_state; - for (int t = 0; t <= sid; t++) - cur_state = (cur_state >> 3) | (((int)outputs[t]) << 6); - } else { - cur_state = ((int)outputs[sid - 2] & 0x7) - | (((int)outputs[sid - 1] & 0x7) << 3) - | (((int)outputs[sid] & 0x7) << 6); - } - acc_t c = d_turbo3_tcq_codebook[cur_state]; - my_recon_sq = c * c; - } - cost[sid] = my_recon_sq; - __syncthreads(); - for (int stride = 256; stride >= 32; stride >>= 1) { - if (sid < stride) cost[sid] += cost[sid + stride]; - __syncthreads(); - } - if (sid < 32) { - acc_t v = cost[sid]; - v += __shfl_down_sync(0xFFFFFFFFULL, v, 16); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 8); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 4); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 2); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 1); - if (sid == 0) cost[0] = v; + // --- recon norm (warp-shuffle sum; lane owns recon positions L, L+32, L+64, L+96) --- + acc_t recon_sq = 0.0f; +#pragma unroll + for (int b = 0; b < 4; b++) { + int pos = lane + 32*b; int cur_state; + if (pos < 2) { cur_state = initial_state; for (int t = 0; t <= pos; t++) cur_state = (cur_state >> 3) | (((int)outputs[t]) << 6); } + else { cur_state = ((int)outputs[pos-2] & 0x7) | (((int)outputs[pos-1] & 0x7) << 3) | (((int)outputs[pos] & 0x7) << 6); } + acc_t cc = d_turbo3_tcq_codebook[cur_state]; recon_sq += cc * cc; } - __syncthreads(); - acc_t recon_norm = tcq_acc_sqrt(cost[0]); +#pragma unroll + for (int o = 16; o > 0; o >>= 1) recon_sq += __shfl_xor_sync(0xFFFFFFFFu, recon_sq, o); + acc_t recon_norm = tcq_acc_sqrt(recon_sq); acc_t corrected_norm = (recon_norm > 1e-10f) ? saved_norm / recon_norm : saved_norm; corrected_norm *= innerq_is_k ? d_tcq_norm_alpha : d_tcq_norm_alpha_v; - // Parallel bitpack: qs stores 6 initial-state bits followed by 128 3-bit - // output symbols. Each byte is independent, so avoid the old serial OR loop. - if (sid < 49) { - const int init_bits = (shared_initial_state >> 3) & 0x3F; - uint8_t packed = 0; -#pragma unroll - for (int bit = 0; bit < 8; bit++) { - const int pos = sid * 8 + bit; - int v = 0; - if (pos < 6) { - v = (init_bits >> pos) & 1; - } else { - const int sym_bit_pos = pos - 6; - const int sym_idx = sym_bit_pos / 3; - if (sym_idx < 128) { - v = (outputs[sym_idx] >> (sym_bit_pos % 3)) & 1; - } + // --- bitpack: lane handles bytes {lane, lane+32} that are < 49 --- + const int init_bits = (initial_state >> 3) & 0x3F; +#pragma unroll + for (int bb = 0; bb < 2; bb++) { + int byte = lane + 32*bb; + if (byte < 49) { + uint8_t packed = 0; +#pragma unroll + for (int bit = 0; bit < 8; bit++) { + const int pos = byte*8 + bit; int v = 0; + if (pos < 6) v = (init_bits >> pos) & 1; + else { const int sym_bit_pos = pos - 6; const int sym_idx = sym_bit_pos / 3; if (sym_idx < 128) v = (outputs[sym_idx] >> (sym_bit_pos % 3)) & 1; } + packed |= (uint8_t)(v << bit); } - packed |= (uint8_t)(v << bit); + dst_blk->qs[byte] = packed; } - dst_blk->qs[sid] = packed; - } - if (sid == 0) { - dst_blk->norm = __float2half((float)corrected_norm); } + if (lane == 0) dst_blk->norm = __float2half((float)corrected_norm); } // TCQ GET_ROWS dequantize (for non-FA paths) @@ -455,10 +347,13 @@ -0.13432936f, -0.05269006f, +0.03536416f, +0.117640756f, -0.022776067f, +0.042032316f, +0.10472976f, +0.18042557f }; -// 2-bit TCQ SET_ROWS encode: Viterbi optimal path with right-shift trellis (k=2, L=8) -// Double-buffered cost arrays + global memory backtrace (128 syncs/group, was 384) +// 2-bit TCQ SET_ROWS encode: warp-synchronous Viterbi (#533, k=2, L=8, 256 states). Same warp redesign as +// turbo3_tcq: 1 group / 32-thread warp (8 states/lane, thread-local quartet predecessor-min, __syncwarp +// broadcast, bt in global), TCQ2_WPB warps/block. Replaces the old 256-thread/__syncthreads block (latency- +// bound on the 128-step barrier chain). bt always global (use_shared_bt ignored); acc_t fp64 path preserved. +#define TCQ2_WPB 8 template -static __global__ void __launch_bounds__(256, 1) k_set_rows_turbo2_tcq( +static __global__ void __launch_bounds__(TCQ2_WPB*32, 3) k_set_rows_turbo2_tcq( const float * __restrict__ src0, const idx_t * __restrict__ src1, block_turbo2_tcq * __restrict__ dst, const int64_t ne_total_groups, uint8_t * __restrict__ bt_buf, @@ -472,12 +367,14 @@ const uint3 ne00_fd, const uint3 ne01_fd, const uint3 ne02_fd, const uint3 ne11_fd, const uint3 ne12_fd) { - const int grp = blockIdx.x; - if (grp >= ne_total_groups) return; - const int sid = threadIdx.x; // 0..255 = trellis state + const int warp = threadIdx.x >> 5; + const int lane = threadIdx.x & 31; + const int64_t group = (int64_t)blockIdx.x * TCQ2_WPB + warp; + if (group >= ne_total_groups) return; + (void) use_shared_bt; (void) ne10; (void) ne13; // bt is always global in the warp design - // Compute source and destination pointers (all threads, used by thread 0) - const int64_t i_base = int64_t(grp) * QK_TURBO2_TCQ; + // Compute source and destination pointers (same index math as the 256-thread design; per-warp group) + const int64_t i_base = group * QK_TURBO2_TCQ; uint32_t tmp = (uint32_t)i_base; uint2 div_mod; div_mod = fast_div_modulo(tmp, ne00_fd); const int64_t i00 = div_mod.y; tmp = div_mod.x; div_mod = fast_div_modulo(tmp, ne01_fd); const int64_t i01 = div_mod.y; tmp = div_mod.x; @@ -489,232 +386,120 @@ block_turbo2_tcq * dst_blk = (block_turbo2_tcq *)((char *)dst + dst_row*s1 + i02*s2 + i03*s3) + (i00 / QK_TURBO2_TCQ); - // Backtrace: one predecessor byte per 64 low-state groups per step. - // The predecessor depends only on sid's low 6 bits (same as turbo3_tcq). - extern __shared__ uint8_t bt_shared[]; - __shared__ acc_t x[128]; - __shared__ acc_t cost[256]; - __shared__ acc_t cost_b[256]; // double-buffering for Viterbi - __shared__ int warp_min_idx[8]; - __shared__ acc_t warp_min_cost[8]; - __shared__ acc_t pred_min_cost[64]; - __shared__ int shared_initial_state; - - if (sid < 128) x[sid] = grp_src[sid]; - __syncthreads(); - - __syncthreads(); - - // Norm reduction - cost[sid] = (sid < 128) ? x[sid] * x[sid] : 0.0f; - __syncthreads(); - for (int stride = 128; stride >= 32; stride >>= 1) { - if (sid < stride) cost[sid] += cost[sid + stride]; - __syncthreads(); - } - if (sid < 32) { - acc_t v = cost[sid]; - v += __shfl_down_sync(0xFFFFFFFFULL, v, 16); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 8); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 4); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 2); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 1); - if (sid == 0) cost[0] = v; - } - __syncthreads(); - acc_t grp_norm = tcq_acc_sqrt(cost[0]); - acc_t inv_norm = grp_norm > 1e-10f ? 1.0f / grp_norm : 0.0f; - - if (sid < 128) x[sid] *= inv_norm; - __syncthreads(); - - // FWHT. The first five stages use warp shuffles, the two cross-warp stages - // use shared memory (same approach as turbo3_tcq). - if (sid < 128) { - acc_t v = x[sid] * TURBO3_WHT_SIGNS1[sid]; - const int lane = sid & 31; -#pragma unroll - for (int h = 1; h < 32; h <<= 1) { - const acc_t other = __shfl_xor_sync(0xFFFFFFFFULL, v, h); - v = (lane & h) ? (other - v) : (v + other); - } - x[sid] = v; - } - __syncthreads(); - if (sid < 64) { - const int j = ((sid >> 5) << 6) + (sid & 31); - acc_t a = x[j], b = x[j + 32]; - x[j] = a + b; x[j + 32] = a - b; - } - __syncthreads(); - if (sid < 64) { - acc_t a = x[sid], b = x[sid + 64]; - x[sid] = a + b; x[sid + 64] = a - b; + // Per-warp smem: whitened vec x[128] (reused as outputs[] after Viterbi) + pred_min_cost[64] broadcast + + // initial-state scratch. Path-cost in registers (c[8], 8 states/lane). bt is global. + __shared__ acc_t x_s[TCQ2_WPB][128]; + __shared__ acc_t pm_s[TCQ2_WPB][64]; + __shared__ int ini_s[TCQ2_WPB]; + acc_t * x = x_s[warp]; + acc_t * pred_min_cost = pm_s[warp]; + uint8_t * bt = bt_buf + group * (int64_t)(128 * 64); // global per-group backtrace slab + + // --- load 4 elements/lane (stride-32 layout: lane owns positions L, L+32, L+64, L+96) --- + acc_t e0 = grp_src[lane], e1 = grp_src[lane+32], e2 = grp_src[lane+64], e3 = grp_src[lane+96]; + // --- norm (warp-shuffle sum of squares) --- + acc_t ss = e0*e0 + e1*e1 + e2*e2 + e3*e3; +#pragma unroll + for (int o = 16; o > 0; o >>= 1) ss += __shfl_xor_sync(0xFFFFFFFFu, ss, o); + acc_t grp_norm = tcq_acc_sqrt(ss); + acc_t inv_norm = grp_norm > 1e-10f ? (acc_t)1.0 / grp_norm : (acc_t)0.0; + e0 *= inv_norm; e1 *= inv_norm; e2 *= inv_norm; e3 *= inv_norm; + // --- FWHT-128 (stride-32 layout: 5 shuffle-butterfly stages + 2 intra-thread cross stages) --- + e0 *= TURBO3_WHT_SIGNS1[lane]; e1 *= TURBO3_WHT_SIGNS1[lane+32]; + e2 *= TURBO3_WHT_SIGNS1[lane+64]; e3 *= TURBO3_WHT_SIGNS1[lane+96]; +#pragma unroll + for (int h = 1; h < 32; h <<= 1) { + acc_t o0 = __shfl_xor_sync(0xFFFFFFFFu, e0, h); e0 = (lane & h) ? o0 - e0 : e0 + o0; + acc_t o1 = __shfl_xor_sync(0xFFFFFFFFu, e1, h); e1 = (lane & h) ? o1 - e1 : e1 + o1; + acc_t o2 = __shfl_xor_sync(0xFFFFFFFFu, e2, h); e2 = (lane & h) ? o2 - e2 : e2 + o2; + acc_t o3 = __shfl_xor_sync(0xFFFFFFFFu, e3, h); e3 = (lane & h) ? o3 - e3 : e3 + o3; } - __syncthreads(); + { acc_t a = e0, b = e1; e0 = a + b; e1 = a - b; acc_t c = e2, d = e3; e2 = c + d; e3 = c - d; } // cross-32 + { acc_t a = e0, b = e2; e0 = a + b; e2 = a - b; acc_t c = e1, d = e3; e1 = c + d; e3 = c - d; } // cross-64 constexpr float inv_sqrt_128 = 0.08838834764831845f; - if (sid < 128) x[sid] *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[sid]; - __syncthreads(); - - if (sid == 0) cost[0] = grp_norm; - __syncthreads(); - - acc_t saved_norm = cost[0]; - - // Viterbi forward pass: double-buffered cost (1 sync/step, was 3) - uint8_t * bt = use_shared_bt ? bt_shared : bt_buf + (int64_t)blockIdx.x * (128 * 64); - cost[sid] = 0.0f; - __syncthreads(); - + e0 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane]; e1 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane+32]; + e2 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane+64]; e3 *= inv_sqrt_128 * TURBO3_WHT_SIGNS2[lane+96]; + x[lane] = e0; x[lane+32] = e1; x[lane+64] = e2; x[lane+96] = e3; + __syncwarp(); + acc_t saved_norm = grp_norm; + + // --- Viterbi forward DP: warp-synchronous (8 states/lane, thread-local quartet pred-min, __syncwarp). --- + // lane L owns states {8L..8L+7} = trellis-groups 2L (states 8L..8L+3) and 2L+1 (states 8L+4..8L+7); each + // group's 4 predecessors (g*4..g*4+3) fall within lane g/2, so predecessor-min is thread-local (no shuffle). + acc_t c[8]; +#pragma unroll + for (int i = 0; i < 8; i++) c[i] = 0.0f; for (int t = 0; t < 128; t++) { - acc_t * cost_rd = (t & 1) ? cost_b : cost; - acc_t * cost_wr = (t & 1) ? cost : cost_b; - + acc_t pm0 = c[0]; int pp0 = 0; +#pragma unroll + for (int j = 1; j < 4; j++) { if (c[j] < pm0) { pm0 = c[j]; pp0 = j; } } + acc_t pm1 = c[4]; int pp1 = 0; +#pragma unroll + for (int j = 1; j < 4; j++) { if (c[4+j] < pm1) { pm1 = c[4+j]; pp1 = j; } } + pred_min_cost[2*lane] = pm0; pred_min_cost[2*lane+1] = pm1; + bt[t*64 + 2*lane] = (uint8_t) pp0; bt[t*64 + 2*lane+1] = (uint8_t) pp1; + __syncwarp(); acc_t xt = x[t]; - - // Right-shift trellis (k=2, L=8): ns = (prev >> 2) | (out << 6). - // The best predecessor depends only on sid's low 6 bits, so compute - // those 64 minima once instead of repeating the same 4-way scan per state. - if (sid < 64) { - const int base_prev = sid << 2; - acc_t best = cost_rd[base_prev]; - int best_p = 0; -#pragma unroll - for (int p = 1; p < 4; p++) { - acc_t c = cost_rd[base_prev | p]; - if (c < best) { - best = c; - best_p = p; - } - } - pred_min_cost[sid] = best; - bt[t * 64 + sid] = (uint8_t) best_p; - } - __syncthreads(); - - const int pred_idx = sid & 0x3F; - acc_t dist = xt - d_turbo2_tcq_codebook[sid]; - dist = dist * dist; - - cost_wr[sid] = pred_min_cost[pred_idx] + dist; - __syncthreads(); - } - // After 128 steps (even count): final costs in cost[] - - // Warp argmin over 256 costs - { - acc_t my_cost = cost[sid]; - int my_idx = sid; - #pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - acc_t other_cost = __shfl_xor_sync(0xFFFFFFFFULL, my_cost, offset); - int other_idx = __shfl_xor_sync(0xFFFFFFFFULL, my_idx, offset); - if (other_cost < my_cost) { my_cost = other_cost; my_idx = other_idx; } - } - if (sid % 32 == 0) { - warp_min_cost[sid / 32] = my_cost; - warp_min_idx[sid / 32] = my_idx; - } +#pragma unroll + for (int i = 0; i < 8; i++) { int s = 8*lane + i; acc_t d = xt - d_turbo2_tcq_codebook[s]; c[i] = pred_min_cost[s & 0x3F] + d * d; } + __syncwarp(); } - __syncthreads(); - if (sid < 32) { - acc_t best = (sid < 8) ? warp_min_cost[sid] : 3.4028234663852886e38f; - int best_idx = (sid < 8) ? warp_min_idx[sid] : 0; -#pragma unroll - for (int offset = 16; offset > 0; offset >>= 1) { - acc_t other_cost = __shfl_down_sync(0xFFFFFFFFULL, best, offset); - int other_idx = __shfl_down_sync(0xFFFFFFFFULL, best_idx, offset); - if (other_cost < best) { - best = other_cost; - best_idx = other_idx; - } - } - if (sid == 0) { - shared_initial_state = best_idx; - } + // --- argmin over 256 states (per-lane min over 8, then warp-shuffle) --- + acc_t mc = c[0]; int mi = 8*lane; +#pragma unroll + for (int i = 1; i < 8; i++) { if (c[i] < mc) { mc = c[i]; mi = 8*lane + i; } } +#pragma unroll + for (int o = 16; o > 0; o >>= 1) { + acc_t oc = __shfl_xor_sync(0xFFFFFFFFu, mc, o); int oi = __shfl_xor_sync(0xFFFFFFFFu, mi, o); + if (oc < mc || (oc == mc && oi < mi)) { mc = oc; mi = oi; } } - __syncthreads(); - - // Save x[] to global buffer before backtrack overwrites it - - // Backtrack (inherently sequential, reads compressed bt) + // --- backtrack (lane 0; outputs reuse the x slab) — turbo2: 2-bit output (state>>6), <<2 trellis shift --- uint8_t * outputs = (uint8_t *)x; - if (sid == 0) { - int state = shared_initial_state; + if (lane == 0) { + int state = mi; for (int t = 127; t >= 0; t--) { outputs[t] = (uint8_t)(state >> 6); - int p = bt[t * 64 + (state & 0x3F)]; + int p = bt[t*64 + (state & 0x3F)]; state = ((state & 0x3F) << 2) | p; } - shared_initial_state = state; + ini_s[warp] = state; } - __syncthreads(); - - // Save output symbols to global buffer + __syncwarp(); + int initial_state = ini_s[warp]; - // Parallel recon norm: t>=3 can compute state directly from 4 outputs (4 shifts of 2 = 8 bits) - acc_t my_recon_sq = 0.0f; - if (sid < 128) { - int cur_state; - if (sid < 3) { - cur_state = shared_initial_state; - for (int t = 0; t <= sid; t++) - cur_state = (cur_state >> 2) | (((int)outputs[t]) << 6); - } else { - cur_state = ((int)outputs[sid - 3] & 0x3) - | (((int)outputs[sid - 2] & 0x3) << 2) - | (((int)outputs[sid - 1] & 0x3) << 4) - | (((int)outputs[sid] & 0x3) << 6); - } - acc_t c = d_turbo2_tcq_codebook[cur_state]; - my_recon_sq = c * c; - } - cost[sid] = my_recon_sq; - __syncthreads(); - for (int stride = 128; stride >= 32; stride >>= 1) { - if (sid < stride) cost[sid] += cost[sid + stride]; - __syncthreads(); - } - if (sid < 32) { - acc_t v = cost[sid]; - v += __shfl_down_sync(0xFFFFFFFFULL, v, 16); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 8); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 4); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 2); - v += __shfl_down_sync(0xFFFFFFFFULL, v, 1); - if (sid == 0) cost[0] = v; + // --- recon norm (warp-shuffle sum; turbo2: pos<3 walk, else 4 symbols x 2-bit) --- + acc_t recon_sq = 0.0f; +#pragma unroll + for (int b = 0; b < 4; b++) { + int pos = lane + 32*b; int cur_state; + if (pos < 3) { cur_state = initial_state; for (int t = 0; t <= pos; t++) cur_state = (cur_state >> 2) | (((int)outputs[t]) << 6); } + else { cur_state = ((int)outputs[pos-3] & 0x3) | (((int)outputs[pos-2] & 0x3) << 2) | (((int)outputs[pos-1] & 0x3) << 4) | (((int)outputs[pos] & 0x3) << 6); } + acc_t cc = d_turbo2_tcq_codebook[cur_state]; recon_sq += cc * cc; } - __syncthreads(); - acc_t recon_norm = tcq_acc_sqrt(cost[0]); +#pragma unroll + for (int o = 16; o > 0; o >>= 1) recon_sq += __shfl_xor_sync(0xFFFFFFFFu, recon_sq, o); + acc_t recon_norm = tcq_acc_sqrt(recon_sq); acc_t corrected_norm = (recon_norm > 1e-10f) ? saved_norm / recon_norm : saved_norm; corrected_norm *= iq_is_k ? d_tcq_norm_alpha : d_tcq_norm_alpha_v; - // Parallel bitpack: qs stores 6 initial-state bits followed by 128 2-bit - // output symbols. Each byte is independent (2-bit symbols never cross byte - // boundaries after the 6-bit prefix), so avoid the serial OR loop. - if (sid < 33) { - const int init_bits = (shared_initial_state >> 2) & 0x3F; - uint8_t packed = 0; -#pragma unroll - for (int bit = 0; bit < 8; bit++) { - const int pos = sid * 8 + bit; - int v = 0; - if (pos < 6) { - v = (init_bits >> pos) & 1; - } else { - const int sym_bit_pos = pos - 6; - const int sym_idx = sym_bit_pos / 2; - if (sym_idx < 128) { - v = (outputs[sym_idx] >> (sym_bit_pos % 2)) & 1; - } + // --- bitpack: 6 init bits + 128 x 2-bit symbols = 33 bytes; lane handles bytes {lane, lane+32} < 33 --- + const int init_bits = (initial_state >> 2) & 0x3F; +#pragma unroll + for (int bb = 0; bb < 2; bb++) { + int byte = lane + 32*bb; + if (byte < 33) { + uint8_t packed = 0; +#pragma unroll + for (int bit = 0; bit < 8; bit++) { + const int pos = byte*8 + bit; int v = 0; + if (pos < 6) v = (init_bits >> pos) & 1; + else { const int sym_bit_pos = pos - 6; const int sym_idx = sym_bit_pos / 2; if (sym_idx < 128) v = (outputs[sym_idx] >> (sym_bit_pos % 2)) & 1; } + packed |= (uint8_t)(v << bit); } - packed |= (uint8_t)(v << bit); + dst_blk->qs[byte] = packed; } - dst_blk->qs[sid] = packed; - } - if (sid == 0) { - dst_blk->norm = __float2half((float)corrected_norm); } + if (lane == 0) dst_blk->norm = __float2half((float)corrected_norm); } // 2-bit TCQ GET_ROWS dequantize