hrx: IQ2_XXS and IQ2_XS weights on HRX0 (shared dequantizer, K-quant decode kernels)

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-kquant-decode.cpp

@@ -14,7 +14,7 @@
 // limitations under the License.
 
 // Decode projections (1 token, or 2-8 for MTP / speculative verify batches) on Q2_K, Q3_K, Q4_K, Q5_K,
-// Q6_K, IQ3_S, IQ4_NL, IQ4_XS and Q8_0 weights, read in their GGUF block layout
+// Q6_K, IQ2_XXS, IQ2_XS, IQ3_S, IQ4_NL, IQ4_XS and Q8_0 weights, read in their GGUF block layout
 // (ops/kquant_decode_f32.loom): FFN gate/up pairs fused with SwiGLU, and plain projections with an
 // optional following residual ADD. Mixed-quant models (Unsloth UD-Q4_K_XL and similar) pair these
 // types freely per layer; without this they take the generic dequantize-4-values-at-a-time
@@ -46,6 +46,8 @@ static constexpr KernelCatalogRef kKQuantMulMatDecodeTokensKernel =
 bool kquant_format(CommonMulMatWeightFormat format) {
     switch (format) {
         case CommonMulMatWeightFormat::Q2K:
+        case CommonMulMatWeightFormat::IQ2_XXS:
+        case CommonMulMatWeightFormat::IQ2_XS:
         case CommonMulMatWeightFormat::Q3K:
         case CommonMulMatWeightFormat::Q4K:
         case CommonMulMatWeightFormat::Q5K:
@@ -60,6 +62,13 @@ bool kquant_format(CommonMulMatWeightFormat format) {
     }
 }
 
+// Formats whose lane functions read a codebook staged in workgroup memory. A gate/up pair shares
+// that one buffer, so its two formats may not need different codebooks.
+bool kquant_needs_grid(CommonMulMatWeightFormat format) {
+    return format == CommonMulMatWeightFormat::IQ3_S || format == CommonMulMatWeightFormat::IQ2_XXS ||
+           format == CommonMulMatWeightFormat::IQ2_XS;
+}
+
 // The kernels' config ranges (ops/kquant_decode_f32.loom).
 constexpr int64_t kMaxInputSize  = 65536;
 constexpr int64_t kMaxOutputSize = 1048576;
@@ -151,6 +160,10 @@ bool match_kquant_swiglu_decode(const DispatchMatchContext & context, DispatchMa
         return false;
     }
     const CommonMulMatMatch other = match_few_token_mul_mat(graph, peer, kKQuantSwiGLUDecodeKernel);
+    if (other.matched() && kquant_needs_grid(other.weight_format) && kquant_needs_grid(root.weight_format) &&
+        other.weight_format != root.weight_format) {
+        return false;
+    }
     if (!other.matched() || !kquant_format(other.weight_format) || other.input->id != root.input->id ||
         other.input_size != root.input_size || other.output_size != root.output_size ||
         other.token_count != root.token_count) {

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-common.h

@@ -185,6 +185,10 @@ inline ggml_type common_mul_mat_format_type(CommonMulMatWeightFormat format) {
             return GGML_TYPE_Q5_0;
         case CommonMulMatWeightFormat::Q5_1:
             return GGML_TYPE_Q5_1;
+        case CommonMulMatWeightFormat::IQ2_XXS:
+            return GGML_TYPE_IQ2_XXS;
+        case CommonMulMatWeightFormat::IQ2_XS:
+            return GGML_TYPE_IQ2_XS;
         case CommonMulMatWeightFormat::IQ2_S:
             return GGML_TYPE_IQ2_S;
         case CommonMulMatWeightFormat::IQ3_S:

ggml/src/ggml-hrx/dispatch_registration/common/dispatch-mul-mat-weight-format.h

@@ -19,6 +19,8 @@ enum class CommonMulMatWeightFormat {
     Q4_1,
     Q5_0,
     Q5_1,
+    IQ2_XXS,
+    IQ2_XS,
     IQ2_S,
     IQ3_S,
     IQ4_NL,
@@ -62,6 +64,12 @@ inline bool common_mul_mat_format_for_type(ggml_type type, CommonMulMatWeightFor
         case GGML_TYPE_Q5_1:
             format = CommonMulMatWeightFormat::Q5_1;
             return true;
+        case GGML_TYPE_IQ2_XXS:
+            format = CommonMulMatWeightFormat::IQ2_XXS;
+            return true;
+        case GGML_TYPE_IQ2_XS:
+            format = CommonMulMatWeightFormat::IQ2_XS;
+            return true;
         case GGML_TYPE_IQ2_S:
             format = CommonMulMatWeightFormat::IQ2_S;
             return true;
@@ -120,6 +128,10 @@ inline int64_t common_mul_mat_format_config_value(CommonMulMatWeightFormat forma
             return 50;
         case CommonMulMatWeightFormat::Q5_1:
             return 51;
+        case CommonMulMatWeightFormat::IQ2_XXS:
+            return 24;
+        case CommonMulMatWeightFormat::IQ2_XS:
+            return 25;
         case CommonMulMatWeightFormat::IQ2_S:
             return 22;
         case CommonMulMatWeightFormat::IQ3_S: