1
0

matrix_view.cuh 5.4 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149
  1. /*
  2. * Adapted from https://github.com/turboderp/exllamav2
  3. * Copyright (c) 2024 turboderp
  4. *
  5. * Permission is hereby granted, free of charge, to any person obtaining a copy
  6. * of this software and associated documentation files (the "Software"), to deal
  7. * in the Software without restriction, including without limitation the rights
  8. * to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
  9. * copies of the Software, and to permit persons to whom the Software is
  10. * furnished to do so, subject to the following conditions:
  11. *
  12. * The above copyright notice and this permission notice shall be included in all
  13. * copies or substantial portions of the Software.
  14. *
  15. * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
  16. * IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
  17. * FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
  18. * AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
  19. * LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
  20. * OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
  21. * SOFTWARE.
  22. */
  23. #ifndef _matrix_view_cuh
  24. #define _matrix_view_cuh
  25. #include <cuda_runtime.h>
  26. #include <cuda_fp16.h>
  27. #include "quant/qdq_util.cuh"
  28. namespace aphrodite {
  29. namespace exl2 {
  30. class MatrixView_half
  31. {
  32. public:
  33. const half* data;
  34. const int height;
  35. const int width;
  36. __device__ __forceinline__ MatrixView_half(const half* data, const int height, const int width)
  37. : data(data), height(height), width(width)
  38. { }
  39. __device__ __forceinline__ half item(int row, int column) const { return data[row * width + column]; }
  40. __device__ __forceinline__ half2 item_half2(int row, int column) const { return ((half2*)data)[(row * width + column) / 2]; }
  41. __device__ __forceinline__ half2 item_half2half2(int row, int column) const { return __half2half2(data[row * width + column]); }
  42. __device__ __forceinline__ const half* item_ptr(int row, int column) const { return &data[row * width + column]; }
  43. __device__ __forceinline__ void item4(half (&items)[4], int row, int column) const
  44. {
  45. half2* ptr = (half2*) item_ptr(row, column);
  46. half2 i01 = ptr[0];
  47. half2 i23 = ptr[1];
  48. items[0] = __low2half(i01);
  49. items[1] = __high2half(i01);
  50. items[2] = __low2half(i23);
  51. items[3] = __high2half(i23);
  52. }
  53. __device__ __forceinline__ void item4_f(float (&items)[4], int row, int column) const
  54. {
  55. half2* ptr = (half2*)item_ptr(row, column);
  56. half2 i01 = ptr[0];
  57. half2 i23 = ptr[1];
  58. items[0] = __half2float(__low2half(i01));
  59. items[1] = __half2float(__high2half(i01));
  60. items[2] = __half2float(__low2half(i23));
  61. items[3] = __half2float(__high2half(i23));
  62. }
  63. __device__ __forceinline__ void item4_h2(half2 (&items)[4], int row, int column) const
  64. {
  65. half2* ptr = (half2*)item_ptr(row, column);
  66. half2 i01 = ptr[0];
  67. half2 i23 = ptr[1];
  68. items[0] = __half2half2(__low2half(i01));
  69. items[1] = __half2half2(__high2half(i01));
  70. items[2] = __half2half2(__low2half(i23));
  71. items[3] = __half2half2(__high2half(i23));
  72. }
  73. };
  74. class MatrixView_half_rw
  75. {
  76. public:
  77. half* data;
  78. const int height;
  79. const int width;
  80. __device__ __forceinline__ MatrixView_half_rw(half* data, const int height, const int width)
  81. : data(data), height(height), width(width)
  82. { }
  83. __device__ __forceinline__ half item(int row, int column) const { return data[row * width + column]; }
  84. __device__ __forceinline__ half2 item_half2(int row, int column) const { return ((half2*)data)[(row * width + column) / 2]; }
  85. __device__ __forceinline__ half* item_ptr(int row, int column) { return &data[row * width + column]; }
  86. __device__ __forceinline__ void set(int row, int column, half value) { data[row * width + column] = value; }
  87. __device__ __forceinline__ void set_half2(int row, int column, half2 value) { ((half2*)data)[(row * width + column) / 2] = value; }
  88. __device__ __forceinline__ void set4(int row, int column, half v0, half v1, half v2, half v3)
  89. {
  90. half2 v01 = __halves2half2(v0, v1);
  91. half2 v23 = __halves2half2(v2, v3);
  92. half2* ptr = (half2*) item_ptr(row, column);
  93. ptr[0] = v01;
  94. ptr[1] = v23;
  95. }
  96. };
  97. class MatrixView_q4_row
  98. {
  99. public:
  100. const uint32_t* data;
  101. const int height;
  102. const int width;
  103. __device__ __forceinline__ MatrixView_q4_row(const uint32_t* data, const int height, const int width)
  104. : data(data), height(height), width(width)
  105. { }
  106. __device__ __forceinline__ int item(int row, int column) const
  107. {
  108. int shift = (column & 0x07) * 4;
  109. return (data[row * width / 8 + column / 8] >> shift) & 0x0f;
  110. }
  111. __device__ __forceinline__ void item2(int (&items)[2], int row, int column) const
  112. {
  113. int shift = (column & 0x07) * 4;
  114. uint32_t d = data[row * width / 8 + column / 8] >> shift;
  115. items[0] = d & 0x0f;
  116. items[1] = (d >> 4) & 0x0f;
  117. }
  118. __device__ __forceinline__ void item4(int (&items)[4], int row, int column) const
  119. {
  120. int shift = (column & 0x07) * 4;
  121. uint32_t d = data[row * width / 8 + column / 8] >> shift;
  122. items[0] = d & 0x0f;
  123. items[1] = (d >> 4) & 0x0f;
  124. items[2] = (d >> 8) & 0x0f;
  125. items[3] = (d >> 12) & 0x0f;
  126. }
  127. };
  128. } // namespace exl2
  129. } // namespace aphrodite
  130. #endif