revise documentation

Author: JohannesGaessler

ggml/src/ggml-cuda/vecdotq.cuh

@@ -28,6 +28,9 @@ static __device__ __forceinline__ int get_int_b4(const void * x, const int & i32
     return ((const int *) x)[i32]; // assume at least 4 byte alignment
 }
 
+// q4 contains 8 indices with 4 bit each.
+// This function selects those bytes from table that are at those indices and returns them as int2.
+// The first int contains the bytes with even indices in q4, the second int contains the bytes with odd indices in q4.
 static __device__ __forceinline__ int2 get_int_from_table_16(const int & q4, const int8_t * table) {
 #if defined(GGML_USE_HIP)
     // Load the 16-byte table into four 32-bit unsigned integers.
@@ -51,25 +54,32 @@ static __device__ __forceinline__ int2 get_int_from_table_16(const int & q4, con
     uint32_t res_y = __builtin_amdgcn_perm(v_odd_high, v_odd_low, mask_odd);
 
     return make_int2(res_x, res_y);
-#elif defined(__CUDA_ARCH__)
-    uint32_t v1, v2, v3, v4, mask;
-    const uint32_t *values = (const uint32_t *)table;
+#elif !defined(GGML_USE_MUSA)
+    // CUDA does not have an instruction for selecting bytes with 4 bit indices.
+    // However, __byte_perm is an instruction that selects bytes with 3 bit indices that can be used instead.
+    const uint32_t * table32 = (const uint32_t *) table;
+
+    // __byte_perm selects bytes based on the lower 16 bits in its third argument.
+    // Therefore, do 2 iterations over the 32 bits in q4 with 0 and 16 shift.
+    // To handle the fourth bit, first call _byte_perm both for the low and the high 64 bit of table, using the low 3 bits.
+    // Then, call __byte_perm again to select from the low and high bytes based on the fourth bit.
+    uint32_t tmp[2];
+    const uint32_t low_high_selection_indices = (0x32103210 | ((q4 & 0x88888888) >> 1));
+#pragma unroll
+    for (uint32_t i = 0; i < 2; ++i) {
+        const uint32_t shift = 16 * i;
 
-    mask = (0x32103210 | ((q4 & 0x88888888) >> 1));
-    // Perform lookups in the lower half of the table (indices 0-7).
-    v1 = __byte_perm(values[0], values[1], q4);
-    // Perform lookups in the upper half of the table (indices 8-15).
-    v2 = __byte_perm(values[2], values[3], q4);
-    // Select between the low and high results based on the MSB of each index nibble.
-    v3 = __byte_perm(v1, v2, mask);
-    // Same for the upper part of q4.
-    v1 = __byte_perm(values[0], values[1], q4 >> 16);
-    v2 = __byte_perm(values[2], values[3], q4 >> 16);
-    v4 = __byte_perm(v1, v2, mask >> 16);
-    
-    // Mix the results to get the final int2.
-    return make_int2(__byte_perm(v3, v4, 0x6420), __byte_perm(v3, v4, 0x7531));
+        const uint32_t low  = __byte_perm(table32[0], table32[1], q4 >> shift);
+        const uint32_t high = __byte_perm(table32[2], table32[3], q4 >> shift);
+        tmp[i] = __byte_perm(low, high, low_high_selection_indices >> shift);
+    }
+
+    // tmp contains the bytes from tyble in the same order as the 4 bit indices in q4.
+    // However, for the result we need ints with all even/odd 4 bit indices in q4.
+    // Therefore, 2 more calls to __byte_perm to put the bytes in the correct order.
+    return make_int2(__byte_perm(tmp[0], tmp[1], 0x6420), __byte_perm(tmp[0], tmp[1], 0x7531));
 #else
+    // Generic implementation.
     const int      q0_32  = (q4 >> 0) & 0x0F0F0F0F;
     const int8_t * q0_8   = (const int8_t *) &q0_32;
     const char4    val0_8 = make_char4(