diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu index f22855a8b..66f38bacf 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/mmvq.cu @@ -167,7 +167,7 @@ static __device__ __forceinline__ float4 vec_dot_rocmfp4_fast_q8_1_x4( for (int l = 0; l < VDR_ROCMFP4_FAST_Q8_1_MMVQ; ++l) { const int aux_q4 = rocmfp4_get_qs_i32(bq4->qs, iqs + l); const int2 v = - rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + rocmfp4_expand_codebook_16(aux_q4); sumi0 = ggml_cuda_dp4a(v.x, q80[l + 0], sumi0); sumi0 = ggml_cuda_dp4a(v.y, q80[l + 4], sumi0); sumi1 = ggml_cuda_dp4a(v.x, q81[l + 0], sumi1); @@ -698,6 +698,57 @@ static constexpr __host__ __device__ int calc_rows_per_block(int ncols_dst, int return 1; } +// Number of independent row groups one CUDA block is split into. +// +// A `small_k` block is dispatched exactly when blocks_per_row_x (= K/qk) is +// below the whole block's blocks_per_iter -- which means the k-loop start +// `tid/(qi/vdr)` already exceeds the block count for the upper warps: they run +// ZERO iterations while the block still produces only `nwarps` rows. On the +// DS4 ROCmFP4-fast decode path that is literally half the block idle +// (`attn_q_b`, K=1024: 32 blocks per row vs blocks_per_iter 64, so warps 2-3 +// are dead). +// +// Splitting the warps into two groups, each covering the full K for its own +// set of `nwarps` rows, doubles the rows a block writes at unchanged +// per-thread work, and halves both the grid and the redundant per-block reads +// of the activation column. +// +// It is bit-exact. Every output row keeps the same per-lane `kbx` assignment +// (group 0's group-local tid IS the old tid; group 1 reproduces exactly what +// the next CUDA block used to compute), and each row's cross-warp reduction +// visits its contributing warps in the same relative order -- the only terms +// dropped are the +0.0f the idle warps used to contribute. +// +// Two groups and not four, even though four (one warp each) would delete the +// cross-warp reduction entirely: that was tried and it FAILS the greedy-output +// sha gate. That failure is measured. The mechanism is only INFERRED (from the +// accumulate structure, not confirmed in SASS): with one warp per group a lane +// runs TWO k-iterations, so the compiler can contract the second accumulate into +// `fma(d, sumi, tmp)` -- one rounding -- where the two-warp form rounds `d*sumi` +// and the cross-warp add separately. If that is right then preserving the +// reduction shape is not sufficient and each warp must also keep running exactly +// one k-iteration; confirm before relying on it (rebuild the four-group variant +// and look for FFMA vs FMUL+FADD in the accumulate -- 3m20, no measure needed). +// +// Scoped to ROCmFP4-fast deliberately: `mmvq.cu` is shared llama.cpp +// infrastructure and this rig's only correctness oracle is a DS4 qtype-101 +// output hash. For a type whose small-K shapes leave fewer warps idle the +// split would reorder the reduction with nothing here able to check it. +static constexpr __host__ __device__ int calc_row_groups( + ggml_type type, int ncols_dst, int table_id, bool small_k, int nwarps) { + return (type == GGML_TYPE_Q4_0_ROCMFP4_FAST && small_k && ncols_dst == 1 && + nwarps % 2 == 0 && + calc_rows_per_block(ncols_dst, table_id, small_k, nwarps) == nwarps) + ? 2 : 1; +} + +// Rows one CUDA block writes: `calc_rows_per_block` for each row group. +static constexpr __host__ __device__ int calc_rows_per_block_total( + ggml_type type, int ncols_dst, int table_id, bool small_k, int nwarps) { + return calc_row_groups(type, ncols_dst, table_id, small_k, nwarps) * + calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); +} + static bool is_gfx1151(const int cc) { return cc == GGML_CUDA_CC_OFFSET_AMD + 0x1151; } @@ -715,7 +766,7 @@ static __device__ __forceinline__ void vec_dot_rocmfp4_fast_q8_1_4cols( #pragma unroll for (int l = 0; l < VDR_ROCMFP4_FAST_Q8_1_MMVQ; ++l) { const int packed = rocmfp4_get_qs_i32(bq4->qs, iqs + l); - weights[l] = rocmfp4_get_int_from_codebook_16(packed, kvalues_rocmfp4); + weights[l] = rocmfp4_expand_codebook_16(packed); } const float weight_scale = rocmfp4_ue4m3_to_fp32_half_finite(bq4->e); @@ -778,11 +829,39 @@ static __global__ void mul_mat_vec_q( constexpr int rows_per_cuda_block = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); constexpr int warp_size = ggml_cuda_get_physical_warp_size(); - const int tid = warp_size*threadIdx.y + threadIdx.x; - const int row0 = rows_per_cuda_block*blockIdx.x; - const int blocks_per_row_x = (fixed_ncols_x > 0 ? fixed_ncols_x : ncols_x) / qk; - constexpr int blocks_per_iter = vdr * nwarps*warp_size / qi; - + // `rows_per_cuda_block` above is the PER ROW GROUP count, so every row loop + // and shared-memory extent below is unchanged; only the block's row origin, + // the group-local thread id and the group's k stride move. + constexpr int row_groups = calc_row_groups(type, ncols_dst, table_id, small_k, nwarps); + constexpr int warps_per_group = nwarps / row_groups; + + const int warp_in_group = threadIdx.y % warps_per_group; + const int row_group = threadIdx.y / warps_per_group; + const int tid = warp_size*warp_in_group + threadIdx.x; + const int row0 = row_groups*rows_per_cuda_block*blockIdx.x + row_group*rows_per_cuda_block; + constexpr int blocks_per_iter = vdr * warps_per_group*warp_size / qi; + // A split block's trailing row group can start past the last row when + // nrows_x is not a multiple of row_groups*rows_per_cuda_block. The result + // write is already guarded; zero the k-loop bound so that group issues no + // out-of-range weight loads either. + const int blocks_per_row_x = + (row_groups > 1 && (uint32_t) row0 >= stride_col_dst) + ? 0 : (int) ((fixed_ncols_x > 0 ? fixed_ncols_x : ncols_x) / qk); + + static_assert(nwarps % row_groups == 0, "row groups must partition the warps"); + static_assert(row_groups == 1 || + (fixed_ncols_x == 0 && !unroll_k_loop_2 && !reuse_rocmfp4_weights && + !c_fp3_packed24 && !c_fp4_x4), + "row-group splitting is limited to the plain small-K k-loop"); + // Freeze the one geometry this was reasoned about: 4 warps -> 2 groups of 2 + // warps, 4 rows each, a group k stride of exactly blocks_per_iter/2, and one + // k-iteration per warp (see above: four groups measurably fails the sha + // gate, most likely via FMA contraction). + static_assert(row_groups == 1 || + (row_groups == 2 && warps_per_group == nwarps/2 && + rows_per_cuda_block == nwarps && + blocks_per_iter == vdr*nwarps*warp_size/(2*qi)), + "unexpected row-group split geometry"); static_assert(fixed_ncols_x == 0 || type == GGML_TYPE_Q3_0_ROCMFPX || type == GGML_TYPE_Q2_0_ROCMFP2, @@ -865,7 +944,7 @@ static __global__ void mul_mat_vec_q( const float * x_bias_base = x_bias + sample_dst*stride_sample_dst + row0; // 1. Hide latency by prefetching bias and gate here // 2. load only on threads that won't die after partial sum calculation - if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 && + if (threadIdx.x < rows_per_cuda_block && warp_in_group == 0 && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { #pragma unroll for (int j = 0; j < ncols_dst; ++j) { @@ -880,7 +959,7 @@ static __global__ void mul_mat_vec_q( } if (use_gate_bias) { const float * gate_bias_base = gate_bias + sample_dst*stride_sample_dst + row0; - if (threadIdx.x < rows_per_cuda_block && threadIdx.y == 0 && + if (threadIdx.x < rows_per_cuda_block && warp_in_group == 0 && (rows_per_cuda_block == 1 || uint32_t(row0 + threadIdx.x) < stride_col_dst)) { #pragma unroll for (int j = 0; j < ncols_dst; ++j) { @@ -1067,24 +1146,24 @@ static __global__ void mul_mat_vec_q( // A one-wave specialization has no cross-wave partials. Avoid reserving // shared memory and issuing a block barrier in that common decode path. // This does not change the per-lane accumulation or warp reduction order. - if constexpr (nwarps > 1) { - __shared__ float tmp_shared[nwarps-1][ncols_dst][rows_per_cuda_block][warp_size]; - __shared__ float tmp_shared_gate[has_fusion ? nwarps-1 : 1][ncols_dst][rows_per_cuda_block][warp_size]; + if constexpr (warps_per_group > 1) { + __shared__ float tmp_shared[row_groups][warps_per_group-1][ncols_dst][rows_per_cuda_block][warp_size]; + __shared__ float tmp_shared_gate[row_groups][has_fusion ? warps_per_group-1 : 1][ncols_dst][rows_per_cuda_block][warp_size]; if constexpr (!has_fusion) { (void) tmp_shared_gate; } else if (!use_gate) { (void) tmp_shared_gate; } - if (threadIdx.y > 0) { + if (warp_in_group > 0) { #pragma unroll for (int j = 0; j < ncols_dst; ++j) { #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { - tmp_shared[threadIdx.y-1][j][i][threadIdx.x] = tmp[j][i]; + tmp_shared[row_group][warp_in_group-1][j][i][threadIdx.x] = tmp[j][i]; if constexpr (has_fusion) { if (use_gate) { - tmp_shared_gate[threadIdx.y-1][j][i][threadIdx.x] = tmp_gate[j][i]; + tmp_shared_gate[row_group][warp_in_group-1][j][i][threadIdx.x] = tmp_gate[j][i]; } } } @@ -1092,7 +1171,7 @@ static __global__ void mul_mat_vec_q( } __syncthreads(); - if (threadIdx.y > 0) { + if (warp_in_group > 0) { return; } @@ -1101,11 +1180,11 @@ static __global__ void mul_mat_vec_q( #pragma unroll for (int i = 0; i < rows_per_cuda_block; ++i) { #pragma unroll - for (int l = 0; l < nwarps-1; ++l) { - tmp[j][i] += tmp_shared[l][j][i][threadIdx.x]; + for (int l = 0; l < warps_per_group-1; ++l) { + tmp[j][i] += tmp_shared[row_group][l][j][i][threadIdx.x]; if constexpr (has_fusion) { if (use_gate) { - tmp_gate[j][i] += tmp_shared_gate[l][j][i][threadIdx.x]; + tmp_gate[j][i] += tmp_shared_gate[row_group][l][j][i][threadIdx.x]; } } } @@ -1792,7 +1871,7 @@ static std::pair calc_launch_params( const int ncols_dst, const int nrows_x, const int nchannels_dst, const int nsamples_or_ntokens, const int warp_size, const mmvq_parameter_table_id table_id, const bool small_k = false) { const int nwarps = calc_nwarps(type, ncols_dst, table_id); - const int rpb = calc_rows_per_block(ncols_dst, table_id, small_k, nwarps); + const int rpb = calc_rows_per_block_total(type, ncols_dst, table_id, small_k, nwarps); const int64_t nblocks = (nrows_x + rpb - 1) / rpb; const dim3 block_nums(nblocks, nchannels_dst, nsamples_or_ntokens); const dim3 block_dims(warp_size, nwarps, 1); diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu index 793485300..f5dc4e33f 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/rocmfp2_mix.cu @@ -326,6 +326,121 @@ __device__ __forceinline__ float mix_fp2_fixed(uint32_t code) { return (float) ((int) code - 1); // 0->-1, 1->0, 2->1, 3->2 } +// ---- nibble-expanded codebook table ---- +// +// Why. After iteration 3's LDS activation staging this fold is no longer memory-bound: +// ncu on the fused MoE launch (H200 sm_90, in=7168 / out=2048 / top-k 4) puts DRAM at +// 11.6% and names ALU the top pipe at 64.4%, with 17.15 M ALU against 10.76 M FMA +// instructions. Hopper has 64 INT32 lanes to 128 FP32 lanes per SM, so at half rate that +// ALU count is ~3x the FMA load: the integer extract/address work AROUND each MAC is the +// cost, not the MAC. Per term the old gather `bk[mix_fp2_code_u64(codes, j)]` spent a +// shift, a mask and an index-scale (3 ALU) plus a 4 B LDS -- 3.67 M times per launch. +// +// The fix. Codes are 2 bits, four to a byte, and a fold consumes them in ascending j, so +// the pair (j, j+1) for EVEN j sits at bit offset 2*(j&3) with (j&3) in {0, 2} -- bit 0 or +// bit 4 of its byte, exactly nibble-aligned. Expand the 4-entry codebook into a 16-entry +// table keyed by that nibble, each entry the float2 {level[n & 3], level[(n >> 2) & 3]}: +// one shift + one mask + one 8 B LDS now serves TWO terms instead of one. +// +// Bit-exact. The table holds the same f32 values the gather read (a codebook level +// widened from bf16 once, or mix_fp2_fixed's exact small integer); each term still +// multiplies by the same scale and the same activation in the same ascending-j order. +// Even j runs 0..30, so a pair lies wholly inside one half of the block and can never +// straddle j == MIX_QK/2 and take the wrong scale. Only the route changes. +// Bank replication, and why the naive table is a LOSS without it. Both the nibble and +// the codebook select are per-lane weight data, so an unreplicated table is a +// data-dependent LDS gather. Laid out flat, the 2 * MIX_NIB = 32 float2 entries span +// 256 B = two bank rounds, so entry e and entry e+16 share a bank pair whatever +// permutation you pick (any 8 B-granular stride shares gcd >= 2 with 32, so 32 entries +// can only reach 16 distinct bank pairs). MEASURED: the flat version cut ALU 17.15 M -> +// 11.43 M exactly as intended, and still ran 51.0 -> 58.4 us on the fused kernel because +// bank conflicts went 0.41 M -> 3.98 M and shared wavefronts 5.61 M -> 9.18 M. +// +// The fix is the standard per-bank LUT replication: MIX_NIB_COPIES = 16 copies of the +// table (one per 8 B bank pair), with copy p of entry e at float2 index +// e * MIX_NIB_COPIES + p, and lane L reading copy L & 15. A 64-bit LDS access is serviced +// in phases of 16 consecutive lanes (16 x 8 B = 128 B = all 32 banks), so every lane of a +// phase reads bank pair (L & 15) -- distinct by construction, for ANY weight data. It +// costs 4 KB of LDS per tensor and a 16-store fill; the fold's address arithmetic is +// unchanged (the entry stride is just 128 B instead of 8 B). +#define MIX_NIB 16 // nibble values per codebook +#define MIX_NIB_ENTRIES (2 * MIX_NIB) // (codebook, nibble) entries +#define MIX_NIB_COPIES 16 // bank-pair copies of each +#define MIX_NIB_SHIFT 7 // log2(MIX_NIB_COPIES * 8 B) +#define MIX_NIB_SLAB (MIX_NIB_ENTRIES * MIX_NIB_COPIES) // float2 per tensor = 4 KB + +// MIX_NIB_SHIFT is hand-computed, and mix_nib_off folds it into a MASK as well as a shift. +// Changing MIX_NIB_COPIES without it would silently read the wrong table entry -- caught by +// the harness hashes, but only if someone runs them. Make it a compile error instead. +static_assert(MIX_NIB_COPIES * sizeof(float2) == (1u << MIX_NIB_SHIFT), + "MIX_NIB_SHIFT must be log2(MIX_NIB_COPIES * sizeof(float2))"); +// 16 copies is not a tunable: a 64-bit LDS phase is 16 lanes, so it is exactly the number +// needed for one copy per bank pair. +static_assert(MIX_NIB_COPIES == 16, "one table copy per 8 B bank pair of a 16-lane phase"); + +// Byte offset of nibble j's entry from the lane's copy base -- nibble(j) * 128 -- in one +// shift and one mask rather than extract-then-scale. `j` is a compile-time constant in +// the fully-unrolled fold, so the shift direction folds away. The highest bit touched is +// 8*(30>>2) + 2*(30&3) + 3 = 63, so the read stays inside `codes`. +__device__ __forceinline__ uint32_t mix_nib_off(uint64_t codes, int j) { + const int bit = 8 * (j >> 2) + 2 * (j & 3); + const uint64_t s = (bit >= MIX_NIB_SHIFT) ? (codes >> (bit - MIX_NIB_SHIFT)) + : (codes << (MIX_NIB_SHIFT - bit)); + return (uint32_t) s & (uint32_t) ((MIX_NIB - 1) << MIX_NIB_SHIFT); +} + +// The two levels for terms (j, j+1), one conflict-free 8 B LDS from the lane's own copy. +__device__ __forceinline__ float2 mix_nib_pair( + const float2 * __restrict__ t, uint64_t codes, int j) { + return *(const float2 *) ((const char *) t + mix_nib_off(codes, j)); +} + +// The lane's copy base within a tensor's slab. Bias the pointer ONCE per kernel rather +// than per gather; MIX_NIB_SLAB is a whole number of copies, so adding a tensor's slab +// offset to an already-biased pointer stays on the same copy. +__device__ __forceinline__ const float2 * mix_nib_lane( + const float2 * __restrict__ slab, int lane) { + return slab + (lane & (MIX_NIB_COPIES - 1)); +} + +// Cooperative fill of one tensor's replicated slab from its resolved-expert codebook. +// Codebook `cb` (selected per block by the meta byte's top bit) owns entries +// [cb * MIX_NIB, (cb+1) * MIX_NIB). mode 0 ignores the codebook entirely -- its levels +// are the fixed {-1, 0, 1, 2} -- and filling BOTH codebooks with them keeps the fold's +// `m >> 7` table select valid there without a second branch in the hot loop. +// +// Written a float4 -- TWO adjacent bank copies of one entry -- per thread per step, so a +// tensor's 4 KB slab costs 4 steps a thread instead of 32. That matters more than it looks: +// the fill is per BLOCK while the fold it feeds is per row-window, so on the dense kernels +// (which have no activation staging to amortise it against) the fill is most of what this +// table costs. +// +// The stride is over the flat float4 index on purpose: consecutive threads write +// consecutive 16 B slots, which is 8 distinct bank quads across the 8 lanes of a 128-bit +// phase -- conflict-free. Grouping it the other way (one thread owns a value and writes +// all MIX_NIB_COPIES copies of it) would save the repeated codebook read but land every +// lane of a store in the SAME bank quad, a 16-way conflict. The repeated read is free: +// `book` is 16 B, so it is an L1 hit after the first. +__device__ __forceinline__ void mix_fill_nib( + float2 * __restrict__ slab, const nv_bfloat16 * __restrict__ book, + int mode, int tid, int nthreads) { + float4 * __restrict__ t4 = (float4 *) slab; + // Two float2 copies per float4 slot, so MIX_NIB_COPIES/2 slots per entry. + const int slots = MIX_NIB_ENTRIES * (MIX_NIB_COPIES / 2); + for (int q = tid; q < slots; q += nthreads) { + const int e = q / (MIX_NIB_COPIES / 2); // (codebook, nibble) entry + const int nib = e & (MIX_NIB - 1); + const int cb = e / MIX_NIB; + const int c0 = nib & 3; // term j + const int c1 = (nib >> 2) & 3; // term j+1 + const float v0 = (mode == 0) ? mix_fp2_fixed((uint32_t) c0) + : __bfloat162float(book[cb * MIX_K + c0]); + const float v1 = (mode == 0) ? mix_fp2_fixed((uint32_t) c1) + : __bfloat162float(book[cb * MIX_K + c1]); + t4[q] = make_float4(v0, v1, v0, v1); + } +} + // One thread per element of a single expert slice (k = out*in elements). __global__ void dequantize_rocmfp2_mix_kernel( const uint8_t * __restrict__ data, const nv_bfloat16 * __restrict__ book, @@ -406,6 +521,42 @@ void dequantize_rocmfp2_mix_to_fp16_cuda(const void * vx, half * y, int64_t k, c #define MIX_WARP 32 #define MIX_UNROLL 8 +// Output rows folded per warp in the MoE matvec (the fused form doubles the FOLDS: +// each row carries an `up` and a `gate` fold). Every row keeps its own accumulator +// and its own fixed ascending-j left fold over the same block sequence, so this is a +// pure work-per-warp knob -- no value of it changes any row's arithmetic or summation +// order, and the output is bit-identical for any choice. +// +// What it trades (measured on H200, decode geometry in=7168/out=2048/n_used=4/ntok=1, +// standalone harness, medians of 7x200 launches, adaptive-codebook mode): +// + FEWER SECTORS. The block's 32 activations are staged ONCE per warp-iteration +// regardless of ROWS, and activations are ~76% of this kernel's global-load +// sectors (lane L owns blocks 32 floats apart, so every float4 activation request +// touches 32 distinct sectors and uses 16 B of each). Doubling ROWS cuts +// sectors/work by 1.76x unfused, 1.6x fused. +// - FEWER WARPS. Total warps = out*n_used*ntok/ROWS. At ROWS=4 that is 2048 warps +// for 132 SMs, and this kernel is latency-bound (66.6% L1, 7.7% DRAM), so +// halving the warp count removes the latency hiding that paid for the sectors. +// +// The two instantiations land on opposite sides of that trade, so they get different +// values rather than one compromise (all bit-identical output): +// unfused (1 fold/row): ROWS 2 -> 4 -> 8 = 47.58 -> **36.13** -> 35.21 us. Its +// sector saving is the larger one. ROWS=8 edges it here but collapses on the +// fixed-codebook mode (38.81 -> 29.95 -> 34.11 us) and leaves only 512 blocks +// = 3.9/SM, so 4. +// fused (2 folds/row): ROWS 2 -> 4 = 66.44 -> **68.96** us, so it keeps 2. It had +// already amortised the stage over 4 folds, so it buys less, and it pays the +// warp-count loss in full. Confirmed concurrency-bound rather than fold-bound by +// an ntok sweep (`bench `): the ROWS=4 penalty shrinks away +// as ntok 1 -> 8 restores the warp count. +// Note ROWS=4 for the fused kernel ALSO fixes the wave quantisation (2048 -> 1024 +// blocks, one wave) and is still slower. Resident warps, not wave count, is the +// objective at ntok=1. +// Must stay in sync with `rows_per_block` in the two mul_mat_id wrappers below. +// The dense and 3-D slice kernels keep their own 2-row blocking and are unaffected. +#define MIX_MOE_ROWS_FUSED 2 +#define MIX_MOE_ROWS_UNFUSED 4 + // Down-shift warp shuffle confined to a 32-lane logical group. width=MIX_WARP // keeps the reduction self-contained on wave64 (GFX8/9, physical wave = 64) and // is a no-op vs the default warp width on wave32 (gfx1151) / NVIDIA, so the @@ -454,6 +605,65 @@ __device__ __forceinline__ void mix_load_x32( } } +// A block's 10 packed bytes, staged in registers: the 32 fp2 codes plus the two +// scale/bank bytes. Split out of mix_block_accum_x so a caller folding several rows +// against one activation stage can issue EVERY row's weight load before any of the +// FMA work (see mix_moe_block_fold). That matters because this kernel is +// latency-bound, not throughput-bound: at the decode geometry a block spends most of +// its time waiting on dependent global loads, so the number of loads in flight is the +// lever. Splitting the phases changes no value and no accumulation order. +struct MixBlock { + uint64_t codes; + uint8_t m0; // block byte MIX_QS+0 + uint8_t m1; // block byte MIX_QS+1 +}; + +__device__ __forceinline__ MixBlock mix_block_load(const uint8_t * __restrict__ b) { + const uintptr_t addr = (uintptr_t) b; + const uint8_t * base8 = (const uint8_t *) MIX_ASSUME_ALIGNED( + (const void *) (addr & ~(uintptr_t) 7), 8); + const int sh = (int) (addr & 7) * 8; // 0, 16, 32, 48 + uint64_t lo, hi; + MIX_MEMCPY(&lo, base8, 8); + MIX_MEMCPY(&hi, base8 + 8, 8); + MixBlock r; + r.codes = (sh == 0) ? lo : ((lo >> sh) | (hi << (64 - sh))); + r.m0 = (uint8_t) (hi >> sh); + r.m1 = (uint8_t) (hi >> (sh + 8)); + return r; +} + +// Fold one already-staged block's 32 terms into acc. Body lifted verbatim out of +// mix_block_accum_x, which is now this plus mix_block_load. +__device__ __forceinline__ void mix_block_fma( + const MixBlock & blk, const float (&xv)[MIX_QK], + int mode, const float2 * __restrict__ nib, float & acc) { + const uint64_t codes = blk.codes; + const uint8_t m0 = blk.m0, m1 = blk.m1; + // mode 0 reads the WHOLE meta byte as the ue4m3 scale; mode 1 reserves its top bit + // for the codebook select, so the scale is the low 7. That difference is the only + // thing left that distinguishes the modes -- the level lookup is now common, because + // mix_fill_nib has already resolved mode 0's fixed {-1, 0, 1, 2} into both tables. + // Folding the two 32-term bodies into one is not just tidier: the mode branch is + // block-uniform, so both unrolled bodies used to be resident in the binary and in the + // register allocator's live-range problem. + const float s0 = mix_ue4m3((mode == 0) ? m0 : (uint8_t) (m0 & 0x7F)); + const float s1 = mix_ue4m3((mode == 0) ? m1 : (uint8_t) (m1 & 0x7F)); + const float2 * t0 = nib + (m0 >> 7) * (MIX_NIB * MIX_NIB_COPIES); + const float2 * t1 = nib + (m1 >> 7) * (MIX_NIB * MIX_NIB_COPIES); + // Two terms per iteration, one nibble and one 8 B LDS each. EVEN j only, so a pair + // never straddles the s0/s1 boundary at j == MIX_QK/2, and the terms still enter acc + // in ascending j -- the summation order the greedy-sha gate depends on is untouched. + #pragma unroll + for (int j = 0; j < MIX_QK; j += 2) { + const float s = (j < MIX_QK/2) ? s0 : s1; + const float2 * t = (j < MIX_QK/2) ? t0 : t1; + const float2 v = mix_nib_pair(t, codes, j); + acc += s * v.x * xv[j]; + acc += s * v.y * xv[j + 1]; + } +} + // Accumulate one block's 32 terms directly into acc, in fixed j order, exactly // as the un-refactored loop did (acc += s * w_j * x_j). Adding each term into // the shared running acc — rather than forming a per-block partial sum first — @@ -464,89 +674,23 @@ __device__ __forceinline__ void mix_load_x32( // compiler overlap their loads even though the acc-add chain stays serial. __device__ __forceinline__ void mix_block_accum_x( const uint8_t * __restrict__ b, const float (&xv)[MIX_QK], - int mode, const float * __restrict__ lut, float & acc) { - // Stage the whole 10-byte block into registers with one 2-byte-wide copy, - // then decode the fp2 codes out of registers instead of re-reading the - // packed qs region through narrow per-byte global loads. Note the load - // pressure differs from qtype-105: at 2 bits each weight touches exactly - // ONE of the 8 qs bytes (four weights share it) rather than up to 3 of 12, - // so staging buys register reuse rather than rescuing a multi-byte gather. - // Blocks are always 2-byte aligned (block stride 10 and row stride - // nb*10=1280 are both even), so a 2-byte assume is safe and lets the - // compiler fold the packed-weight reads into ushort loads. A 10-byte block - // is also 8+2, so a u64+u16 pair is a natural next step -- left for the - // evolution loop to try against the correctness gate rather than assumed. The - // decode arithmetic and the fixed j accumulation order are untouched, so - // acc is bit-for-bit identical to the per-byte path (the correctness gate - // hashes the greedy output; any reassociation flips a token). - // Load-from-floor wide staging (per-lane, no lane->block remap). The block is - // 10 B at byte offset 10*bidx from an 8-aligned rowbase, so its address is - // 2-byte- (not 8-byte-) aligned: (10*bidx) & 7 cycles {0,2,4,6}. Rather than - // the ~5 narrow ushort loads a direct 10-byte copy forces, load the two - // 8-aligned u64 that bracket the block (floor = addr & ~7; the block spans - // [r, r+10) with r<=6, so r+10<=16 always fits in the 16 B window), then - // funnel-shift the lane's OWN 10 bytes out of registers. Same block, same - // bytes, same decode, same fixed j order, same acc-add chain -- only the load - // *route* changes, NOT which block-dots enter this lane's acc, so the fold is - // not reassociated and the greedy-sha gate is unaffected. 2 dwordx2 vs ~5 - // ushort per block. OOB-safe on this model: floor >= rowbase (rowbase is - // 8-aligned). The 16 B window ends at floor+16 = block_end + (6 - (addr&7)); - // for the LAST block of the tensor that overruns the row end by - // (6 - (10*(nb-1) & 7)) B, which is 0 exactly when nb % 4 == 0 (in % 128 == 0). - // in=4096 -> nb=128 -> overrun 0 (last block reads [rowbase+1272, rowbase+1280), - // exactly to the 8-aligned row end). If a future shape has nb % 4 != 0, the last - // block reads up to 6 B past the tensor end. That is now REJECTED AT REGISTRATION - // (see the in % 128 guard in ggml_cuda_rocmfp2_mix_register_host) rather than - // left to chance, so this loop can stay branch-free. - const uintptr_t addr = (uintptr_t) b; - const uint8_t * base8 = (const uint8_t *) MIX_ASSUME_ALIGNED( - (const void *) (addr & ~(uintptr_t) 7), 8); - const int sh = (int) (addr & 7) * 8; // 0, 16, 32, 48 - uint64_t lo, hi; - MIX_MEMCPY(&lo, base8, 8); - MIX_MEMCPY(&hi, base8 + 8, 8); - const uint64_t codes = (sh == 0) ? lo : ((lo >> sh) | (hi << (64 - sh))); - const uint8_t m0 = (uint8_t) (hi >> sh); // block byte MIX_QS+0 - const uint8_t m1 = (uint8_t) (hi >> (sh + 8)); // block byte MIX_QS+1 - if (mode == 0) { - const float s0 = mix_ue4m3(m0), s1 = mix_ue4m3(m1); - #pragma unroll - for (int j = 0; j < MIX_QK; ++j) { - const float s = (j < MIX_QK/2) ? s0 : s1; - acc += s * mix_fp2_fixed(mix_fp2_code_u64(codes, j)) * xv[j]; - } - } else { - // `lut` is the expert's 2 * MIX_K codebook ALREADY widened to f32 and staged in - // LDS by the caller (see the staging blocks in the two matvec kernels). It used - // to be a global bf16 pointer dereferenced inside this unrolled loop, which - // issued MIX_QK dependent 2-byte global loads per 32-weight block for a table - // that is workgroup-invariant and only 16 B -- pure LSU issue and dependent-load - // latency, never bandwidth, since all lanes hit the same bytes. - // - // Bit-exact: __bfloat162float is a widening with no rounding, so hoisting it to - // the staging loop changes no value, and the fold is still - // acc += s * level * x in the same ascending-j order with the same two roundings - // per term. The correctness gate hashes the greedy output, so that matters. - const float s0 = mix_ue4m3(m0 & 0x7F), s1 = mix_ue4m3(m1 & 0x7F); - const float * bk0 = lut + (m0 >> 7) * MIX_K; - const float * bk1 = lut + (m1 >> 7) * MIX_K; - #pragma unroll - for (int j = 0; j < MIX_QK; ++j) { - const float s = (j < MIX_QK/2) ? s0 : s1; - const float * bk = (j < MIX_QK/2) ? bk0 : bk1; - acc += s * bk[mix_fp2_code_u64(codes, j)] * xv[j]; - } - } + int mode, const float2 * __restrict__ nib, float & acc) { + // Exactly mix_block_load + mix_block_fma, and deliberately not a second copy of + // the 32-term fold: the dense/slice kernels and the MoE kernel must not be able + // to drift apart numerically. Both callees are __forceinline__, so this composes + // to the same code the one-piece version emitted (verified: dense 64 regs, slice + // 70, harness hashes unchanged). + mix_block_fma(mix_block_load(b), xv, mode, nib, acc); } // Original (col0-addressed) entry point, kept for the MoE and 3-D slice kernels. // Stages the activations itself; identical arithmetic. __device__ __forceinline__ void mix_block_accum( const uint8_t * __restrict__ b, const float * __restrict__ xc, int col0, - int mode, const float * __restrict__ lut, float & acc) { + int mode, const float2 * __restrict__ nib, float & acc) { float xv[MIX_QK]; mix_load_x32(xc + col0, xv); - mix_block_accum_x(b, xv, mode, lut, acc); + mix_block_accum_x(b, xv, mode, nib, acc); } // Two output rows, ONE activation stage. The 32 activations a block consumes are @@ -561,12 +705,176 @@ __device__ __forceinline__ void mix_block_accum( // for element. __device__ __forceinline__ void mix_block_accum2( const uint8_t * __restrict__ b0, const uint8_t * __restrict__ b1, - const float * __restrict__ xp, int mode, const float * __restrict__ lut, + const float * __restrict__ xp, int mode, const float2 * __restrict__ nib, float & acc0, float & acc1) { float xv[MIX_QK]; mix_load_x32(xp, xv); - mix_block_accum_x(b0, xv, mode, lut, acc0); - mix_block_accum_x(b1, xv, mode, lut, acc1); + mix_block_accum_x(b0, xv, mode, nib, acc0); + mix_block_accum_x(b1, xv, mode, nib, acc1); +} + +// ---- LDS activation staging for the MoE fold ---- +// +// The problem this solves. Lane L owns blocks L, L+MIX_WARP, ..., so lane L consumes +// activations xcol[32*L .. 32*L+31] -- consecutive lanes are 32 floats = 128 B apart. +// A warp-wide float4 activation request therefore touches 32 DISTINCT 128 B lines and +// uses only 16 B of each: 8 requests x 32 lines = 256 L1 line-lookups per +// warp-iteration, which is why ncu put this kernel at 89.7% (pre-iter-1) / 66.6% +// (post) L1/TEX throughput with DRAM at <8%. It is bound on L1 tag-lookup +// wavefronts, not bandwidth. +// +// But the warp's 32 lanes together consume a CONTIGUOUS 1024-float (4 KB) window per +// iteration -- only the lane->element map is a permutation. So load that window +// COALESCED (each request 32 consecutive float4 = 512 B = 4 lines: 8 x 4 = 32 +// line-lookups, an 8x cut) into LDS, then let each lane read its own 32 floats back. +// +// The swizzle. A block's 32 floats are 8 float4 chunks; chunk c of the window's local +// block b lands at word `b*MIX_QK + 4*((c + b) & 7)`. A 128-bit LDS access is serviced +// in phases of 8 lanes (8 x 16 B = 128 B = all 32 banks), so conflict-freedom only +// needs the 8 lanes of a phase to hit disjoint bank quads: +// read-back (b = lane, c fixed): bank = 4*((c + lane) & 7); over any 8 consecutive +// lanes (c + lane) & 7 takes all 8 values -> 8 disjoint quads = all 32 banks. +// fill (b = g>>3, c = g&7, g = lane + i*MIX_WARP): within an 8-lane phase b is fixed +// and c runs 0..7 -> again all 8 values. +// A rotate rather than the 36-float pad the obvious fix would use: same +// conflict-freedom, but 4 KB/warp instead of 4.5 KB and no dead words. +__device__ __forceinline__ int mix_swz_word(int b, int c) { + return b * MIX_QK + 4 * ((c + b) & (MIX_QK / 4 - 1)); +} + +// Stage `nblk` blocks' worth of a contiguous activation window into this warp's LDS +// buffer with coalesced float4 loads. Bit-exact: the SAME f32 values reach the fold, +// only the route changes. +// +// The trip count is exact, never ragged: mix_validate_shape() forces in % 128 == 0, so +// nb % 4 == 0, and `nblk` is either MIX_WARP or nb % MIX_WARP -- a multiple of 4 either +// way -- hence nblk * (MIX_QK/4) is a multiple of MIX_WARP. No lane sits out a +// half-iteration, and the caller's __syncwarp() sees a fully-written buffer. +// +// The alignment test is warp-uniform (the window base is a multiple of MIX_QK floats = +// 128 B, so it inherits xcol's alignment), so the branch costs no divergence. The +// scalar arm is live: ggml can hand this path a strided/offset src1 column. +// FULL says the window is a whole MIX_WARP blocks, which makes `iters` the compile-time +// constant MIX_QK/4 and lets the fill unroll completely. It matters: with a runtime trip +// count nvcc emits a rolled loop with an unroll-remainder epilogue, and the model's shape +// (nb=224 = 7 full windows) would pay that even though it never runs a partial window. +template +__device__ __forceinline__ void mix_stage_window( + float * __restrict__ s, const float * __restrict__ xw, int nblk, int lane) { + const int iters = FULL ? MIX_QK / 4 : nblk * (MIX_QK / 4) / MIX_WARP; + if ((((uintptr_t) xw) & 15u) == 0) { + const float4 * __restrict__ p4 = (const float4 *) xw; + #pragma unroll + for (int i = 0; i < iters; ++i) { + const int g = lane + i * MIX_WARP; + const float4 v = p4[g]; + float * __restrict__ d = s + mix_swz_word(g / (MIX_QK / 4), g & (MIX_QK / 4 - 1)); + d[0] = v.x; d[1] = v.y; d[2] = v.z; d[3] = v.w; + } + } else { + #pragma unroll + for (int i = 0; i < iters; ++i) { + const int g = lane + i * MIX_WARP; + const float * __restrict__ q = xw + 4 * g; + float * __restrict__ d = s + mix_swz_word(g / (MIX_QK / 4), g & (MIX_QK / 4 - 1)); + d[0] = q[0]; d[1] = q[1]; d[2] = q[2]; d[3] = q[3]; + } + } +} + +// Read local block `b`'s 32 activations back out of the staged window, ascending j, as +// 8 conflict-free 128-bit LDS loads. `s` is 16 B aligned and mix_swz_word is a multiple +// of 4 words, so every float4 access is aligned. +__device__ __forceinline__ void mix_read_x32( + const float * __restrict__ s, int b, float (&xv)[MIX_QK]) { + #pragma unroll + for (int c = 0; c < MIX_QK / 4; ++c) { + const float4 v = *(const float4 *) (s + mix_swz_word(b, c)); + xv[4*c + 0] = v.x; xv[4*c + 1] = v.y; + xv[4*c + 2] = v.z; xv[4*c + 3] = v.w; + } +} + +// One block of the MoE fold: stage the block's 32 activations ONCE, then fold them +// against this warp's ROWS `up` rows and (fused only) its ROWS `gate` rows. +// +// Measured motivation (H200 sm_90, ncu on the real decode geometry +// in=7168/out=2048/top-k 4): the FUSED kernel sat at 89.7% L1/TEX throughput with +// DRAM at 3.9% and issue slots at 21% -- bound on L1 tag-lookup wavefronts, not on +// bandwidth, arithmetic or occupancy. 31.7 M global-load sectors against 1.16 M +// requests is 27.2 sectors/request, i.e. essentially the 32-sector worst case, +// because consecutive lanes own blocks 32 floats apart: each float4 activation +// request touches 32 distinct sectors and uses only 16 of each 32 B. +// +// The caller used to invoke mix_block_accum() once per fold, each of which staged the +// SAME 32 activations from the SAME address. The comment there asserted nvcc would CSE +// those loads to one issue per element; the SASS says otherwise (288 static LDG.E.128 +// for a loop that needs 36), so the fused path paid the 32-sector activation request +// four times per block. Staging once cuts activation requests 4x on the fused path and +// 2x on the unfused one. This is the same register-blocking mix_block_accum2() already +// does for the dense 2-D path, extended to the fused gate/up pair. +// +// Bit-exact: identical to the four mix_block_accum() calls it replaces. The same 32 +// activation values arrive from the same addresses, and each row is folded by the same +// mix_block_accum_x into its own accumulator over the same fixed ascending-j order, so +// every acc keeps its exact left-fold summation order. Only the load *route* changes, +// not which terms enter which accumulator in what order -- required, because the +// correctness gate hashes the greedy token stream and any reassociation could flip it. +template +__device__ __forceinline__ void mix_moe_block_fold( + // The row bases are a plain (non-__restrict__) array on purpose: a tail warp + // whose row index is past `out` has that slot clamped onto row0 by the caller, + // so asserting no-alias here would be a lie for that shape. Every base is + // read-only const, so permitting the aliasing costs nothing. + const uint8_t * const (&rowbase)[ROWS], const uint8_t * const (&growbase)[ROWS], + const float (&xv)[MIX_QK], int b, + int mode, int gmode, const float2 * __restrict__ nib, + float (&acc)[ROWS], float (&gacc)[ROWS]) { + const int64_t off = (int64_t) b * MIX_BLOCK_BYTES; + // Issue all 2*ROWS weight-block loads before any FMA. Leaving the loads inside + // per-row accumulate calls let nvcc serialise row i's load behind row i-1's + // 32-term fold, which on a latency-bound kernel is the whole cost. The + // activations arrive pre-staged in `xv` (the caller reads them out of LDS), so + // this block issues nothing to L1 but weights. + MixBlock w[ROWS], gw[ROWS]; + #pragma unroll + for (int i = 0; i < ROWS; ++i) w[i] = mix_block_load(rowbase[i] + off); + if (FUSE_GLU) { + #pragma unroll + for (int i = 0; i < ROWS; ++i) gw[i] = mix_block_load(growbase[i] + off); + } + #pragma unroll + for (int i = 0; i < ROWS; ++i) mix_block_fma(w[i], xv, mode, nib, acc[i]); + if (FUSE_GLU) { + #pragma unroll + for (int i = 0; i < ROWS; ++i) mix_block_fma(gw[i], xv, gmode, nib + MIX_NIB_SLAB, gacc[i]); + } +} + +// One MIX_WARP-block window of the MoE fold: stage this warp's activation window into +// LDS, then fold the lane's own block against it. FULL drops both the runtime trip count +// in the fill and the `lane < nblk` predicate, so the hot path carries neither. +template +__device__ __forceinline__ void mix_moe_window( + float * __restrict__ sxw, const float * __restrict__ xcol, int base, int nblk, int lane, + const uint8_t * const (&rowbase)[ROWS], const uint8_t * const (&growbase)[ROWS], + int mode, int gmode, const float2 * __restrict__ nib, + float (&acc)[ROWS], float (&gacc)[ROWS]) { + mix_stage_window(sxw, xcol + (int64_t) base * MIX_QK, nblk, lane); + // Required, not warp-synchrony folklore: sm_90 schedules threads independently, so a + // lane's LDS writes are not ordered against a sibling lane's reads without it. + // Warp-scoped, so unlike __syncthreads() it does not couple the block's two warps and + // cost the latency hiding iteration 2 showed this kernel lives on. + __syncwarp(); + if (FULL || lane < nblk) { + float xv[MIX_QK]; + mix_read_x32(sxw, lane, xv); + mix_moe_block_fold(rowbase, growbase, xv, base + lane, + mode, gmode, nib, acc, gacc); + } + // WAR: a lane that races ahead to the next window's fill must not overwrite slots a + // sibling lane is still reading. + __syncwarp(); } // The lane's block loop is unrolled by MIX_UNROLL into a SINGLE accumulator kept @@ -589,17 +897,16 @@ __global__ void mix_matvec_rocmfp2_kernel( const int row = warp * 2; // two output rows per warp const int lane = threadIdx.x % MIX_WARP; const int col = blockIdx.y; - // Widen the 2 * MIX_K bf16 codebook to f32 in LDS once per workgroup instead of - // re-loading it from global per weight inside mix_block_accum. HOISTED ABOVE the - // row early-return: __syncthreads() requires every thread of the workgroup to - // arrive, and `row >= out` retires whole warps in the tail block. - __shared__ float s_lut[2 * MIX_K]; - if ((int) threadIdx.x < 2 * MIX_K) { - s_lut[threadIdx.x] = __bfloat162float(book[threadIdx.x]); - } + // Expand the bf16 codebook into the nibble-keyed f32 table (see mix_fill_nib) once + // per workgroup instead of re-loading it from global per weight inside + // mix_block_accum. HOISTED ABOVE the row early-return: __syncthreads() requires every + // thread of the workgroup to arrive, and `row >= out` retires whole warps in the tail + // block. `mode` is hoisted with it because the table's contents depend on it. + const int mode = (int) mode_ptr[0]; + __shared__ __align__(16) float2 s_nib[MIX_NIB_SLAB]; + mix_fill_nib(s_nib, book, mode, (int) threadIdx.x, (int) blockDim.x); __syncthreads(); if (row >= out) return; - const int mode = (int) mode_ptr[0]; const int nb = in / MIX_QK; // `two` is warp-uniform, so the hot loop never diverges. It is false only for // the tail warp of an odd `out`; that warp reuses row0's base so every load @@ -628,14 +935,14 @@ __global__ void mix_matvec_rocmfp2_kernel( const int b = blk + u * MIX_WARP; mix_block_accum2(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, - xc + b * MIX_QK, mode, s_lut, acc0, acc1); + xc + b * MIX_QK, mode, mix_nib_lane(s_nib, lane), acc0, acc1); } } // Remainder: fewer than MIX_UNROLL strided blocks left for this lane. for (; blk < nb; blk += MIX_WARP) { mix_block_accum2(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, - xc + blk * MIX_QK, mode, s_lut, acc0, acc1); + xc + blk * MIX_QK, mode, mix_nib_lane(s_nib, lane), acc0, acc1); } #pragma unroll for (int off = MIX_WARP/2; off > 0; off >>= 1) { @@ -707,16 +1014,16 @@ __global__ void mix_matvec_rocmfp2_slice_kernel( const int expert = slot; // DENSE: the weight slice is the grid slice, // not a routing lookup. This is the ONLY // change from the MoE kernel. - __shared__ float s_lut[2 * MIX_K]; - if ((int) threadIdx.x < 2 * MIX_K) { - s_lut[threadIdx.x] = - __bfloat162float(codebooks[(int64_t) expert * 2 * MIX_K + threadIdx.x]); - } + // `mode` is read up here with the expert because mix_fill_nib's table contents + // depend on it (mode 0's levels are the fixed set, not the codebook). + const int mode = (int) modes[expert]; + __shared__ __align__(16) float2 s_nib[MIX_NIB_SLAB]; + mix_fill_nib(s_nib, codebooks + (int64_t) expert * 2 * MIX_K, mode, + (int) threadIdx.x, (int) blockDim.x); __syncthreads(); if (row0 >= out) return; const bool two = (row0 + 1) < out; // false only for an odd-out tail warp const uint8_t * edata = data + (int64_t) expert * nb02; - const int mode = (int) modes[expert]; const int nb = in / MIX_QK; const uint8_t * rowbase0 = edata + (int64_t) row0 * nb * MIX_BLOCK_BYTES; // For an odd tail warp (no row1) reuse row0's base so the loads stay in-bounds; @@ -740,13 +1047,13 @@ __global__ void mix_matvec_rocmfp2_slice_kernel( #pragma unroll for (int u = 0; u < MIX_UNROLL; ++u) { const int b = blk + u * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc1); + mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, mix_nib_lane(s_nib, lane), acc0); + mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, mix_nib_lane(s_nib, lane), acc1); } } for (; blk < nb; blk += MIX_WARP) { - mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); + mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, mix_nib_lane(s_nib, lane), acc0); + mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, mix_nib_lane(s_nib, lane), acc1); } #pragma unroll for (int off = MIX_WARP/2; off > 0; off >>= 1) { @@ -773,7 +1080,7 @@ __global__ void mix_matvec_rocmfp2_slice_kernel( // Naming follows the mmvq fusion convention: the PRIMARY tensor is `up` (src0 of the surviving // mul_mat_id) and `gate` arrives as the extra operand, because // ggml_cuda_op_swiglu_ds4_single(gate, up, limit) is not symmetric -- silu() is applied to gate. -template +template __global__ void mix_matvec_rocmfp2_moe_kernel( const uint8_t * __restrict__ data, size_t nb02, const nv_bfloat16 * __restrict__ codebooks, const uint8_t * __restrict__ modes, @@ -788,9 +1095,12 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( const uint8_t * __restrict__ gdata, size_t gnb02, const nv_bfloat16 * __restrict__ gcodebooks, const uint8_t * __restrict__ gmodes, float glu_limit) { - const int warps_per_block = blockDim.x / MIX_WARP; - const int warp = blockIdx.x * warps_per_block + (threadIdx.x / MIX_WARP); - const int row0 = warp * 2; // two output rows per warp + // WARPS is a template parameter, not blockDim.x/MIX_WARP, because it also sizes the + // per-warp LDS activation buffer below -- a runtime warp count could index past it. + // Both mul_mat_id wrappers launch WARPS * MIX_WARP threads. + const int wib = threadIdx.x / MIX_WARP; // warp index within the block + const int warp = blockIdx.x * WARPS + wib; + const int row0 = warp * MIX_MOE_ROWS; // MIX_MOE_ROWS output rows per warp const int lane = threadIdx.x % MIX_WARP; const int slot = blockIdx.y; const int token = blockIdx.z; @@ -807,18 +1117,33 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( // an invalid id to a ZERO contribution for this (token, slot): deterministic, and the // same thing a masked-out slot means. All threads still reach the barrier below. const bool bad_expert = expert < 0 || expert >= n_experts; - // Both tables in one array so the staging stays a single guarded write per thread and one - // barrier. Gate's table occupies [2*MIX_K, 4*MIX_K). - __shared__ float s_lut[FUSE_GLU ? 4 * MIX_K : 2 * MIX_K]; - if (!bad_expert && (int) threadIdx.x < 2 * MIX_K) { - s_lut[threadIdx.x] = - __bfloat162float(codebooks[(int64_t) expert * 2 * MIX_K + threadIdx.x]); - } - if (FUSE_GLU) { - const int gt = (int) threadIdx.x - 2 * MIX_K; - if (!bad_expert && gt >= 0 && gt < 2 * MIX_K) { - s_lut[2 * MIX_K + gt] = - __bfloat162float(gcodebooks[(int64_t) expert * 2 * MIX_K + gt]); + // Both tensors' nibble tables in one array so the fill stays one barrier. Gate's table + // occupies [MIX_NIB_SLAB, 2*MIX_NIB_SLAB) -- the offset the fold adds for the gate fold. + // 4 KB per tensor against the old 32 B; 16 KB total per block still leaves this kernel + // register-limited (96 regs -> 10 blocks/SM = 160 KB of the 228 KB an SM has). + __shared__ __align__(16) float2 s_nib[FUSE_GLU ? 2 * MIX_NIB_SLAB : MIX_NIB_SLAB]; + // One MIX_WARP-block activation window per warp (4 KB), swizzled -- see + // mix_swz_word. PER-WARP rather than shared by the whole block on purpose: both + // warps do walk the identical block sequence, so one buffer would serve both, but + // that needs __syncthreads() in the hot loop and this block has only 2 warps. + // Iteration 2 measured that trading warp-level independence for fewer loads loses + // here (MIX_MOE_ROWS=4 on the fused kernel: 1.6x fewer sectors, 0.96x the speed), + // and cross-warp sharing would only take activation line-lookups 32 -> 16 after + // coalescing has already taken them 256 -> 32. Not worth a block barrier. + __shared__ __align__(16) float s_xw[WARPS][MIX_WARP * MIX_QK]; + // The modes are read here rather than after the barrier because mix_fill_nib's table + // contents depend on them (mode 0's levels are the fixed set, not the codebook). A bad + // expert must not index modes/codebooks at all, so it fills a mode-0 table it will + // never read -- every thread still reaches the barrier below. + const int mode = bad_expert ? 0 : (int) modes[expert]; + const int gmode = (FUSE_GLU && !bad_expert) ? (int) gmodes[expert] : 0; + if (!bad_expert) { + mix_fill_nib(s_nib, codebooks + (int64_t) expert * 2 * MIX_K, mode, + (int) threadIdx.x, (int) blockDim.x); + if (FUSE_GLU) { + mix_fill_nib(s_nib + MIX_NIB_SLAB, + gcodebooks + (int64_t) expert * 2 * MIX_K, gmode, + (int) threadIdx.x, (int) blockDim.x); } } __syncthreads(); @@ -826,87 +1151,102 @@ __global__ void mix_matvec_rocmfp2_moe_kernel( if (bad_expert) { if (lane == 0) { const int64_t o = (int64_t) token * dst_s2 + (int64_t) slot * dst_s1 + row0; - dst[o] = 0.0f; - if (row0 + 1 < out) dst[o + 1] = 0.0f; + #pragma unroll + for (int i = 0; i < MIX_MOE_ROWS; ++i) { + if (row0 + i < out) dst[o + i] = 0.0f; + } } return; } - const bool two = (row0 + 1) < out; // false only for an odd-out tail warp const uint8_t * edata = data + (int64_t) expert * nb02; - const int mode = (int) modes[expert]; const int nb = in / MIX_QK; - const uint8_t * rowbase0 = edata + (int64_t) row0 * nb * MIX_BLOCK_BYTES; - // For an odd tail warp (no row1) reuse row0's base so the loads stay in-bounds; - // acc1 is simply never written. out=hidden is even here so `two` is uniformly - // true across the whole warp (no divergence in the hot loop). - const uint8_t * rowbase1 = two ? edata + (int64_t) (row0 + 1) * nb * MIX_BLOCK_BYTES - : rowbase0; // Gate shares the shape, the expert and the row indices -- only the bytes and the table - // differ -- so it reuses `nb`, `two`, `row0` and the same activation column below. - const uint8_t * gedata = FUSE_GLU ? gdata + (int64_t) expert * gnb02 : nullptr; - const int gmode = FUSE_GLU ? (int) gmodes[expert] : 0; - const uint8_t * growbase0 = FUSE_GLU ? gedata + (int64_t) row0 * nb * MIX_BLOCK_BYTES - : nullptr; - const uint8_t * growbase1 = FUSE_GLU ? (two ? gedata + (int64_t) (row0 + 1) * nb * MIX_BLOCK_BYTES - : growbase0) - : nullptr; + // differ -- so it reuses `nb`, `row0` and the same activation column below. + const uint8_t * gedata = FUSE_GLU ? gdata + (int64_t) expert * gnb02 : nullptr; + // Row bases for this warp's MIX_MOE_ROWS rows. A tail warp whose row index is past + // `out` CLAMPS onto row0 so its loads stay in bounds; that accumulator is simply + // never stored. `out` is a multiple of MIX_MOE_ROWS for this model (2048), so the + // clamp is uniform across the warp and costs no divergence in the hot loop -- but it + // must stay correct for a ragged `out`, which neither the model nor the microbench + // exercises. + const uint8_t * rowbase[MIX_MOE_ROWS]; + const uint8_t * growbase[MIX_MOE_ROWS]; + #pragma unroll + for (int i = 0; i < MIX_MOE_ROWS; ++i) { + const int r = (row0 + i) < out ? row0 + i : row0; + rowbase[i] = edata + (int64_t) r * nb * MIX_BLOCK_BYTES; + growbase[i] = FUSE_GLU ? gedata + (int64_t) r * nb * MIX_BLOCK_BYTES : nullptr; + } // src1 is [in, ne11, ntok]; the get_rows-equivalent row for (slot, token) // is token*ne11 + slot%ne11 — i.e. token column + the slot%ne11 broadcast. const float * xcol = src1 + (int64_t) token * src1_s2 + (int64_t) (slot % ne11) * src1_s1; - // Two output rows in one warp. Each row is folded by the SAME mix_block_accum + // Two output rows in one warp. Each row is folded by the SAME mix_block_accum_x // that the single-row path uses (byte-identical inlined body, same fixed j // order, same acc-add chain) so acc0/acc1 are bit-for-bit identical to the - // single-row kernel's output for those rows. The two calls per block share the - // same __restrict__ xcol + col0, so the compiler CSEs the strided activation - // loads to one issue per element — halving activation LSU issue on this partly - // load-instruction-bound matvec — WITHOUT reordering either row's summation. - float acc0 = 0.0f, acc1 = 0.0f; + // single-row kernel's output for those rows. The folds share ONE explicit + // activation stage per block (mix_moe_block_fold) instead of relying on the + // compiler to CSE a separate stage per fold: it does not (the SASS showed 288 + // static LDG.E.128 for a loop that needs 36), and on this L1-wavefront-bound + // kernel that stage is the most expensive request in the loop. Sharing it does + // NOT reorder either row's summation. + float acc[MIX_MOE_ROWS]; // Gate accumulates in its own registers over the SAME block order, so gacc is bit-identical - // to what the separate gate launch produced. All four folds share one `xcol`, so the fused - // form reads the activation column once for four rows instead of once for two -- on a launch - // that is partly load-issue bound, that is the second saving after the launch itself. - float gacc0 = 0.0f, gacc1 = 0.0f; - int blk = lane; - for (; blk + (MIX_UNROLL - 1) * MIX_WARP < nb; blk += MIX_UNROLL * MIX_WARP) { - #pragma unroll - for (int u = 0; u < MIX_UNROLL; ++u) { - const int b = blk + u * MIX_WARP; - mix_block_accum(rowbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) b * MIX_BLOCK_BYTES, xcol, b * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } - } + // to what the separate gate launch produced. ALL folds consume ONE staged copy of the + // block's activations, so the fused form reads the activation column once for + // 2*MIX_MOE_ROWS rows instead of once per row -- on a launch measured at 89.7% L1/TEX + // throughput with DRAM at 3.9%, that is the second saving after the launch itself, and + // the larger of the two. + float gacc[MIX_MOE_ROWS]; + #pragma unroll + for (int i = 0; i < MIX_MOE_ROWS; ++i) { acc[i] = 0.0f; gacc[i] = 0.0f; } + // Walk the row in MIX_WARP-block windows. Lane L still folds blocks L, L+MIX_WARP, + // ... in ascending order into its own accumulator, exactly as the strided loop this + // replaces did -- so every acc keeps its bit-identical left-fold summation order. + // What changed is only how the activations reach `xv`: staged once per window, + // coalesced, through LDS instead of 8 32-line-wide global requests per block. + // + // Windowing (rather than the old `blk += MIX_WARP` stride) is what makes the + // staging legal at all: with nb % MIX_WARP != 0 the strided loop lets lanes exit at + // different trip counts, so a departed lane would stop contributing to the + // cooperative fill while its neighbours read slots nobody wrote. + // + // This also retires the MIX_UNROLL path, which was dead code: its guard + // `blk + 7*MIX_WARP < nb` is false on the first test at the model's nb=224, so every + // block already went through the tail loop. Ordering is unaffected either way -- + // the unrolled body visited the same ascending b sequence. + float * __restrict__ sxw = s_xw[wib]; + const int nfull = nb / MIX_WARP; + for (int w = 0; w < nfull; ++w) { + mix_moe_window( + sxw, xcol, w * MIX_WARP, MIX_WARP, lane, + rowbase, growbase, mode, gmode, mix_nib_lane(s_nib, lane), acc, gacc); } - for (; blk < nb; blk += MIX_WARP) { - mix_block_accum(rowbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc0); - mix_block_accum(rowbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, mode, s_lut, acc1); - if (FUSE_GLU) { - mix_block_accum(growbase0 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc0); - mix_block_accum(growbase1 + (int64_t) blk * MIX_BLOCK_BYTES, xcol, blk * MIX_QK, gmode, s_lut + 2 * MIX_K, gacc1); - } + // At most ONE partial window, and only for shapes with nb % MIX_WARP != 0 (the model's + // nb=224 has none). Its block count is a multiple of 4 -- mix_validate_shape forces + // in % 128 == 0 -- so the fill's trip count stays exact there too. + if (nb % MIX_WARP) { + mix_moe_window( + sxw, xcol, nfull * MIX_WARP, nb % MIX_WARP, lane, + rowbase, growbase, mode, gmode, mix_nib_lane(s_nib, lane), acc, gacc); } #pragma unroll for (int off = MIX_WARP/2; off > 0; off >>= 1) { - acc0 += mix_warp_shfl_down(acc0, off); - acc1 += mix_warp_shfl_down(acc1, off); - if (FUSE_GLU) { - gacc0 += mix_warp_shfl_down(gacc0, off); - gacc1 += mix_warp_shfl_down(gacc1, off); + #pragma unroll + for (int i = 0; i < MIX_MOE_ROWS; ++i) { + acc[i] += mix_warp_shfl_down(acc[i], off); + if (FUSE_GLU) gacc[i] += mix_warp_shfl_down(gacc[i], off); } } if (lane == 0) { const int64_t o = (int64_t) token * dst_s2 + (int64_t) slot * dst_s1 + row0; - if (FUSE_GLU) { - // The SAME function the standalone swiglu_ds4 kernel applies, on inputs bit-identical - // to the ones it would have read back from the two intermediates. - dst[o] = ggml_cuda_op_swiglu_ds4_single(gacc0, acc0, glu_limit); - if (two) dst[o + 1] = ggml_cuda_op_swiglu_ds4_single(gacc1, acc1, glu_limit); - } else { - dst[o] = acc0; - if (two) dst[o + 1] = acc1; + #pragma unroll + for (int i = 0; i < MIX_MOE_ROWS; ++i) { + if (row0 + i >= out) continue; // ragged-`out` tail row: clamped, not owned + // The SAME function the standalone swiglu_ds4 kernel applies, on inputs + // bit-identical to the ones it would have read back from the two intermediates. + dst[o + i] = FUSE_GLU ? ggml_cuda_op_swiglu_ds4_single(gacc[i], acc[i], glu_limit) + : acc[i]; } } } @@ -973,11 +1313,11 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id( } const int warps_per_block = 2; // 64 threads (mirror the mmvq path) const int threads = warps_per_block * MIX_WARP; - // Two output rows per warp (register-blocked activation reuse), so a workgroup - // of `warps_per_block` warps covers 2*warps_per_block rows. - const int rows_per_block = 2 * warps_per_block; + // MIX_MOE_ROWS_UNFUSED output rows per warp (register-blocked activation reuse), + // so a workgroup of `warps_per_block` warps covers that many times as many rows. + const int rows_per_block = MIX_MOE_ROWS_UNFUSED * warps_per_block; dim3 grid((out + rows_per_block - 1) / rows_per_block, n_expert_used, n_tokens); - mix_matvec_rocmfp2_moe_kernel<<>>( + mix_matvec_rocmfp2_moe_kernel<<>>( (const uint8_t *) e.base, e.nb02, e.codebooks, e.modes, src1, ids, dst, in, out, e.n_experts, ne11, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, @@ -1015,9 +1355,9 @@ bool ggml_cuda_rocmfp2_mix_mul_mat_id_glu( } const int warps_per_block = 2; const int threads = warps_per_block * MIX_WARP; - const int rows_per_block = 2 * warps_per_block; + const int rows_per_block = MIX_MOE_ROWS_FUSED * warps_per_block; dim3 grid((out + rows_per_block - 1) / rows_per_block, n_expert_used, n_tokens); - mix_matvec_rocmfp2_moe_kernel<<>>( + mix_matvec_rocmfp2_moe_kernel<<>>( (const uint8_t *) eu.base, eu.nb02, eu.codebooks, eu.modes, src1, ids, dst, in, out, eu.n_experts, ne11, ids_s0, ids_s1, src1_s1, src1_s2, dst_s1, dst_s2, diff --git a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh index f39e8cf5f..7ac12e549 100644 --- a/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh +++ b/server/deps/llama.cpp/ggml/src/ggml-cuda/vecdotq.cuh @@ -467,6 +467,115 @@ static __device__ __forceinline__ int rocmfpx_pack4_fp2_vec_cuda(const uint8_t * return rocmfpx_pack4_fp2_bits8_vec_cuda(qs[base >> 2]); } +// Expand eight packed ROCmFP4 nibbles into two DP4A operands, table-free. +// +// The shared helper `rocmfp4_get_int_from_codebook_16` only has a permute-based +// fast path under HIP; every other target falls through to a 16-entry table +// expander, and on CUDA `kvalues_rocmfp4` is a `__device__` global, so each +// call compiles to EIGHT data-dependent `LDG.E.U8` gathers (measured in SASS: +// 8 gathers / ~44 instructions). In batch-1 MMVQ that lookup, not the weight +// stream, is the throughput limit. +// +// CUDA has `prmt` (`__byte_perm`), which llama.cpp already uses for the same +// job in `get_int_from_table_16` above, so reuse that structure here with the +// ROCmFP4 codebook inlined as four immediates instead of loaded from the table +// (0 gathers / ~19 instructions). Unlike AMD's V_PERM_B32, prmt takes one +// selector NIBBLE per output byte from the low 16 bits of its third argument, +// hence the two 16-bit halves and the final even/odd re-gather. +// +// Bit-exactness is not an approximation argument: the four immediates are the +// bytes of `kvalues_rocmfp4` verbatim, and the mapping was verified equal to +// the table expander for all 2^32 inputs of `q4` on sm_90. Identical integer +// operands feed the identical DP4A order, so the float result is unchanged. +static __device__ __forceinline__ int2 rocmfp4_expand_codebook_16(const int & q4) { +#if defined(GGML_USE_HIP) + // AMD already has a permute path with the codebook inlined; leave it alone. + return rocmfp4_get_int_from_codebook_16(q4, kvalues_rocmfp4); +#elif !defined(GGML_USE_MUSA) + constexpr uint32_t values0 = 0x03020100u; // 0, 1, 2, 3 + constexpr uint32_t values1 = 0x0a080604u; // 4, 6, 8, 10 + constexpr uint32_t values2 = 0xfdfeff00u; // 0, -1, -2, -3 + constexpr uint32_t values3 = 0xf6f8fafcu; // -4, -6, -8,-10 + + const uint32_t q4u = (uint32_t) q4; + // Low 3 bits of each nibble pick the magnitude, the 4th bit picks the sign + // half of the codebook by biasing the selector into the `high` operand. + const uint32_t low_high_selection_indices = 0x32103210u | ((q4u & 0x88888888u) >> 1); + + uint32_t tmp[2]; +#pragma unroll + for (uint32_t i = 0; i < 2; ++i) { + const uint32_t shift = 16*i; + const uint32_t low = __byte_perm(values0, values1, q4u >> shift); + const uint32_t high = __byte_perm(values2, values3, q4u >> shift); + tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift); + } + + // tmp holds the codebook bytes in nibble order; split into even/odd streams. + return make_int2((int) __byte_perm(tmp[0], tmp[1], 0x6420), + (int) __byte_perm(tmp[0], tmp[1], 0x7531)); +#else + return rocmfp4_get_int_from_codebook_16(q4, kvalues_rocmfp4); +#endif +} + +// Load `nints` CONSECUTIVE int32 of a ROCmFP4 `qs` array with aligned dword +// loads instead of 4*nints byte loads. +// +// `block_rocmfp4_fast` is 17 bytes (16 qs + 1 scale), so only one block in four +// starts 4-byte aligned and the portable branch of `rocmfp4_get_qs_i32` +// assembles every int32 from four separate `LDG.E.U8`. At VDR=2 that is eight +// byte loads per `vec_dot`; post-codebook-inlining they are essentially all the +// memory traffic the kernel issues (measured in SASS: 9 of 9 `LDG.E.U8`, the +// ninth being the scale). CUDA cannot issue an unaligned `ld.global.u32` (the +// HIP branch above can, which is why AMD has no problem here), but 4*nints +// consecutive bytes span only nints+1 aligned dwords, and `__funnelshift_r` +// re-assembles each int32 from an adjacent pair. 8 loads -> 3. +// +// Bit-exact: the bytes delivered are the same bytes, in the same order, so the +// DP4A operands and their summation order are unchanged. +// +// Bounds. With `r = addr & 3` the reads cover [addr-r, addr-r+4*nints+3]. +// - Below `addr`: r > 0 requires a preceding block in the same tensor (block 0 +// of a tensor is 128-byte aligned => r == 0), so the low end never precedes +// the tensor. +// - Above: the final dword is issued only when r != 0, and then it reaches at +// most 2 bytes past the last block. That is still in bounds, in both of the +// only two cases `ggml_backend_cuda_buffer_type_get_alloc_size()` allows: +// ne0 % MATRIX_ROW_PADDING(512) == 0, so ne0/32 is a multiple of 16, so the +// tensor's block count is a multiple of 4, so its last block has r == 3 and +// the reach is exactly the scale byte; or ne0 % 512 != 0, in which case that +// function appends >= one zeroed block (17 B) of padding after ggml_nbytes. +template +static __device__ __forceinline__ void rocmfp4_load_qs_i32x( + const uint8_t * __restrict__ qs, const int & i32, int (&out)[nints]) { +#if defined(GGML_USE_HIP) || defined(GGML_USE_MUSA) +#pragma unroll + for (int i = 0; i < nints; ++i) { + out[i] = rocmfp4_get_qs_i32(qs, i32 + i); + } +#else + const uintptr_t addr = (uintptr_t) (qs + 4*i32); + const uint32_t r = (uint32_t) (addr & 3); + const uint32_t shift = 8*r; + const uint32_t * __restrict__ p = (const uint32_t *) (addr - r); + + uint32_t d[nints + 1]; +#pragma unroll + for (int i = 0; i < nints; ++i) { + d[i] = p[i]; + } + // Only the low r bytes of the last dword are consumed, so skip it entirely + // when r == 0 -- that is the one case where it would read past the block. + d[nints] = r != 0 ? p[nints] : 0; + +#pragma unroll + for (int i = 0; i < nints; ++i) { + out[i] = (int) __funnelshift_r(d[i], d[i + 1], shift); + } +#endif +} + static __device__ __forceinline__ float vec_dot_rocmfp4_q8_1( const void * __restrict__ vbq, const block_q8_1 * __restrict__ bq8_1, const int & kbx, const int & iqs) { @@ -479,7 +588,7 @@ static __device__ __forceinline__ float vec_dot_rocmfp4_q8_1( #pragma unroll for (int l = 0; l < VDR_ROCMFP4_Q8_1_MMVQ; ++l) { const int aux_q4 = rocmfp4_get_qs_i32(bq4->qs, iqs + l); - const int2 v = rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + const int2 v = rocmfp4_expand_codebook_16(aux_q4); sumi0 = ggml_cuda_dp4a(v.x, q8[l + 0], sumi0); sumi1 = ggml_cuda_dp4a(v.y, q8[l + 4], sumi1); @@ -496,11 +605,13 @@ static __device__ __forceinline__ float vec_dot_rocmfp4_fast_q8_1( const int * q8 = (const int *) bq8_1->qs + iqs; + int aux_q4[VDR_ROCMFP4_FAST_Q8_1_MMVQ]; + rocmfp4_load_qs_i32x(bq4->qs, iqs, aux_q4); + int sumi = 0; #pragma unroll for (int l = 0; l < VDR_ROCMFP4_FAST_Q8_1_MMVQ; ++l) { - const int aux_q4 = rocmfp4_get_qs_i32(bq4->qs, iqs + l); - const int2 v = rocmfp4_get_int_from_codebook_16(aux_q4, kvalues_rocmfp4); + const int2 v = rocmfp4_expand_codebook_16(aux_q4[l]); sumi = ggml_cuda_dp4a(v.x, q8[l + 0], sumi); sumi = ggml_cuda_dp4a(v.y, q8[l + 4], sumi);