feat(ml8): native 4-bit token_embd via GGML_OP_ML8_GET_ROWS + role-uniform tiering

Adds a native ml8-4 token embedding path so token_embd can be 4-bit (centroid
LUT in an F8_E4M3 sidecar) instead of fp8 — reclaiming size on the largest
tensor. The generic get_rows aborts on ML8_4 (no LUT slot), so this introduces
a dedicated custom op:

  GGML_OP_ML8_GET_ROWS  (ggml_ml8_get_rows(w, centroids, ids))

- ggml op def + name/symbol tables + asserts (op count 102->103)
- CPU compute: LUT-dequant gather (ggml-cpu/ops.cpp) + dispatch + n_tasks
- CUDA/HIP kernel + host dispatch (ml8.cu/.cuh) — pure gather, no AITER needed;
  + ggml-cuda.cu compute dispatch + supports_op
- loader op-swap GET_ROWS -> ML8_GET_ROWS for ML8_4 token_embd (keeps it on GPU)
  + synthetic probe case
- build_inp_embd routes ML8_4 token_embd through the new op via the registry
- qwen35: register token_embd centroids; tied LM head shares token_embd's LUT
  (no separate output.centroids in a tied GGUF)

Role-uniform tier override (ML8_TIER_OVERRIDE / role_targets) lets a role be
moved uniformly between ml8(4-bit) and fp8(8-bit) — the "tread shaping". The
qwen35 FFN fused path now dispatches up/down per-tensor: ml8_4 stays inline
(byte-identical), a role-uniform fp8 tier routes through build_lora_mm. The
data-free ml8-4 embed calibration pass + converter emit token_embd as ML8_4.

Validated on dense Qwen3.5-0.8B (token_embd=ml8, ssm_out/ffn_down/attn_v=fp8,
--awq none): PPL 19.7179 / 498MB — beats the uniform-ml8 clean q3 (19.9387 /
558MB) on BOTH axes; smaller than UD-Q4_K_XL (546MB) but +0.78 PPL behind.

Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>

Author: kmbandy

ggml/include/ggml-ml8.h

@@ -102,6 +102,27 @@ GGML_API struct ggml_tensor * ggml_ml8_mul_mat_id(
         struct ggml_tensor  * x,
         struct ggml_tensor  * ids);
 
+// Native ml8-4 row gather (token-embedding lookup). Gathers row `ids[i]` from the
+// ml8-4 weight `w` [K, N] and dequantizes it via the per-K-group centroid LUT to
+// K fp32 values — kept native 4-bit (no inline bf16 dequant of the table).
+//
+// Tensor shapes (ggml row-major):
+//   w         : [K, N]            GGML_TYPE_ML8_4   (N = vocab rows)
+//   centroids : [16, n_groups_k]  GGML_TYPE_F8_E4M3 per-K-group LUT
+//   ids       : [n_rows, ...]     GGML_TYPE_I32
+//
+// Output:
+//   y         : [K, n_rows, ...]  GGML_TYPE_F32     (matches ggml_get_rows layout)
+//
+// Constraints:
+//   - K % QK_ML8 == 0
+//   - centroids ne0 == 16, ne1 == K / QK_ML8
+GGML_API struct ggml_tensor * ggml_ml8_get_rows(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * w,
+        struct ggml_tensor  * centroids,
+        struct ggml_tensor  * ids);
+
 #ifdef __cplusplus
 }  // extern "C"
 #endif

ggml/include/ggml.h

@@ -625,6 +625,15 @@ extern "C" {
         // HIP backend dispatches to mt_ml8_moe_gemm; CPU fallback dequantizes
         // each routed expert block-by-block.
         GGML_OP_ML8_MUL_MAT_ID,
+        // ml8-4 native row gather (token-embedding lookup for an ml8-4
+        // token_embd). Unlike GGML_OP_GET_ROWS, the row data is ML8_4 and is
+        // dequantized via the per-K-group centroid LUT sidecar — kept native
+        // 4-bit (no inline bf16 dequant of the whole table).
+        //   src[0] = w (GGML_TYPE_ML8_4, [K, N])  — embedding rows (N = vocab)
+        //   src[1] = centroids (GGML_TYPE_F8_E4M3, [16, K/QK_ML8])
+        //   src[2] = ids (GGML_TYPE_I32, [n_rows, ...])
+        //   dst    = y (GGML_TYPE_F32, [K, n_rows, ...])
+        GGML_OP_ML8_GET_ROWS,
 
         GGML_OP_COUNT,
     };

ggml/src/ggml-cpu/ggml-cpu.c

@@ -2164,6 +2164,11 @@ static void ggml_compute_forward(struct ggml_compute_params * params, struct ggm
                 ggml_compute_forward_ml8_mul_mat_id(params, tensor);
             }
             break;
+        case GGML_OP_ML8_GET_ROWS:
+            {
+                ggml_compute_forward_ml8_get_rows(params, tensor);
+            }
+            break;
         case GGML_OP_NONE:
             {
                 // nop
@@ -2507,6 +2512,7 @@ static int ggml_get_n_tasks(struct ggml_tensor * node, int n_threads) {
         case GGML_OP_ML8_MUL_MAT:
         case GGML_OP_ML8_APPLY_ROTATION:
         case GGML_OP_ML8_MUL_MAT_ID:
+        case GGML_OP_ML8_GET_ROWS:
             {
                 n_tasks = n_threads;
             } break;

ggml/src/ggml-cpu/ops.cpp

@@ -11480,6 +11480,67 @@ void ggml_compute_forward_ml8_mul_mat(const ggml_compute_params * params, ggml_t
     free(w_row_fp32);
 }
 
+// ─── ggml_compute_forward_ml8_get_rows (MAD-256 ml8-4 native embed) ───────
+//
+// CPU dispatch for GGML_OP_ML8_GET_ROWS — native 4-bit token-embedding gather.
+// For each requested row id, dequantize that ml8-4 weight row directly into the
+// output via dequantize_row_ml8_4_with_lut (the per-K-group centroid LUT
+// sidecar). The table stays native 4-bit — no inline bf16 dequant of W. Same
+// addressing as ggml_compute_forward_get_rows_q, tiled across gathered rows.
+//
+// Tensor layouts (set by ggml_ml8_get_rows):
+//   dst->src[0] = w         (GGML_TYPE_ML8_4 ,  [K, N]  N = vocab rows)
+//   dst->src[1] = centroids (GGML_TYPE_F8_E4M3, [16, K/QK_ML8] shared LUT)
+//   dst->src[2] = ids       (GGML_TYPE_I32)
+//   dst         = y         (GGML_TYPE_F32   ,  [K, ids->ne0, ids->ne1, ids->ne2])
+void ggml_compute_forward_ml8_get_rows(const ggml_compute_params * params, ggml_tensor * dst) {
+    const ggml_tensor * w   = dst->src[0];
+    const ggml_tensor * lut = dst->src[1];
+    const ggml_tensor * ids = dst->src[2];
+    GGML_ASSERT(w && lut && ids);
+    GGML_ASSERT(w->type   == GGML_TYPE_ML8_4);
+    GGML_ASSERT(lut->type == GGML_TYPE_F8_E4M3);
+    GGML_ASSERT(ids->type == GGML_TYPE_I32);
+    GGML_ASSERT(dst->type == GGML_TYPE_F32);
+
+    const int64_t K = w->ne[0];
+    const int64_t N = w->ne[1];
+    GGML_ASSERT(dst->ne[0] == K);
+    GGML_ASSERT(K % QK_ML8 == 0);
+    const int64_t n_groups_k = K / QK_ML8;
+    GGML_ASSERT(lut->ne[0] == 16 && lut->ne[1] == n_groups_k);
+
+    const uint8_t * lut_fp8 = (const uint8_t *) lut->data;
+
+    // gathered rows: flatten over all ids elements (mirrors get_rows_q)
+    const int64_t nr = ggml_nelements(ids);
+    GGML_ASSERT(ggml_nrows(dst) == nr);
+
+    const int64_t ne10 = ids->ne[0];
+    const int64_t ne11 = ids->ne[1];
+
+    const int ith = params->ith;
+    const int nth = params->nth;
+    const int64_t dr  = (nr + nth - 1) / nth;
+    const int64_t ir0 = dr * ith;
+    const int64_t ir1 = MIN(ir0 + dr, nr);
+
+    for (int64_t i = ir0; i < ir1; i++) {
+        const int64_t i12 = i / (ne11 * ne10);
+        const int64_t i11 = (i - i12 * ne11 * ne10) / ne10;
+        const int64_t i10 = (i - i12 * ne11 * ne10 - i11 * ne10);
+        const int64_t row = *(const int32_t *) ((const char *) ids->data
+                              + i10 * ids->nb[0] + i11 * ids->nb[1] + i12 * ids->nb[2]);
+        GGML_ASSERT(row >= 0 && row < N);
+
+        const block_ml8_4 * w_row =
+            (const block_ml8_4 *) ((const char *) w->data + row * w->nb[1]);
+        float * y_row = (float *) ((char *) dst->data
+                          + i10 * dst->nb[1] + i11 * dst->nb[2] + i12 * dst->nb[3]);
+        dequantize_row_ml8_4_with_lut(w_row, lut_fp8, y_row, K);
+    }
+}
+
 // ─── ggml_compute_forward_ml8_apply_rotation (MAD-223 G.4.g) ──────────────
 //
 // Y[:, t] = unflatten(H_a^T @ flatten(X[:, t]) @ H_b, d=a*b)

ggml/src/ggml-cpu/ops.h

@@ -118,6 +118,7 @@ void ggml_compute_forward_opt_step_sgd(const struct ggml_compute_params * params
 void ggml_compute_forward_ml8_mul_mat(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_ml8_apply_rotation(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 void ggml_compute_forward_ml8_mul_mat_id(const struct ggml_compute_params * params, struct ggml_tensor * dst);
+void ggml_compute_forward_ml8_get_rows(const struct ggml_compute_params * params, struct ggml_tensor * dst);
 #ifdef __cplusplus
 }
 #endif

ggml/src/ggml-cuda/ggml-cuda.cu

@@ -3261,6 +3261,9 @@ static bool ggml_cuda_compute_forward(ggml_backend_cuda_context & ctx, struct gg
         case GGML_OP_ML8_MUL_MAT_ID:
             ggml_cuda_op_ml8_mul_mat_id(ctx, dst);
             break;
+        case GGML_OP_ML8_GET_ROWS:
+            ggml_cuda_op_ml8_get_rows(ctx, dst);
+            break;
         case GGML_OP_OUT_PROD:
             ggml_cuda_out_prod(ctx, dst);
             break;
@@ -5582,6 +5585,24 @@ static bool ggml_backend_cuda_device_supports_op(ggml_backend_dev_t dev, const g
                 if (w->ne[2] <= 0)                    return false;
                 return true;
             } break;
+        case GGML_OP_ML8_GET_ROWS:
+            {
+                // MAD-256 ml8-4 native token-embedding gather. Pure dequant
+                // gather (no AITER GEMM) — ml8_4 weight, f8_e4m3 centroid LUT,
+                // i32 ids, fp32 output. K must be a multiple of QK_ML8=64.
+                const ggml_tensor * w    = op->src[0];
+                const ggml_tensor * cent = op->src[1];
+                const ggml_tensor * ids  = op->src[2];
+                if (!w || !cent || !ids) return false;
+                if (w->type    != GGML_TYPE_ML8_4)   return false;
+                if (cent->type != GGML_TYPE_F8_E4M3) return false;
+                if (ids->type  != GGML_TYPE_I32)     return false;
+                if (op->type   != GGML_TYPE_F32)     return false;
+                if (w->ne[0] % 64 != 0)              return false;
+                if (cent->ne[0] != 16)               return false;
+                if (cent->ne[1] != w->ne[0] / 64)    return false;
+                return true;
+            } break;
         case GGML_OP_OUT_PROD:
             return op->type == GGML_TYPE_F32 && op->src[0]->type == GGML_TYPE_F32 && op->src[1]->type == GGML_TYPE_F32;
         case GGML_OP_GET_ROWS:

ggml/src/ggml-cuda/ml8.cu

@@ -1012,6 +1012,94 @@ void ggml_cuda_op_ml8_mul_mat(
 #endif // GGML_HIP_AITER
 }
 
+// ─────────────────────────────────────────────────────────────────────
+// GGML_OP_ML8_GET_ROWS — native 4-bit token-embedding gather.
+//
+// One CUDA block per gathered row; threads stride over the row's K-groups.
+// For each group: read the per-block fp32 scale + 32 packed nibbles from the
+// native block_ml8_4 layout, index the shared per-group centroid LUT (16 fp8
+// e4m3 each), dequant = centroid * scale, write K fp32. No AITER GEMM, no
+// repack — this is a pure gather so it works on any CUDA/HIP build. Mirrors
+// the CPU ggml_compute_forward_ml8_get_rows math exactly (same LUT, same
+// lo-nibble-first ordering, same e4m3→fp32 helper).
+// ─────────────────────────────────────────────────────────────────────
+static __global__ void ml8_get_rows_kernel(
+    const block_ml8_4 * __restrict__ w,    // [N rows][n_groups_k blocks] native layout
+    const uint8_t     * __restrict__ lut,  // [n_groups_k, 16] fp8 e4m3 (flat g*16+i)
+    const int32_t     * __restrict__ ids,  // [nr] contiguous
+    float             * __restrict__ y,    // [nr, K] row-major (K contiguous per row)
+    int K, int N, int n_groups_k, int64_t nr) {
+
+    const int64_t i = blockIdx.x;          // gathered-row index
+    if (i >= nr) return;
+
+    const int32_t row = ids[i];
+    // out-of-range ids would read garbage rows; clamp defensively to 0.
+    const int32_t row_safe = (row >= 0 && row < N) ? row : 0;
+
+    const block_ml8_4 * w_row = w + (int64_t) row_safe * n_groups_k;
+    float             * y_row = y + i * (int64_t) K;
+
+    for (int g = threadIdx.x; g < n_groups_k; g += blockDim.x) {
+        const block_ml8_4 * blk   = &w_row[g];
+        const float         scale = blk->scale;
+        const uint8_t     * lut_g = lut + (int64_t) g * 16;
+        const int           k_base = g * QK_ML8;
+        #pragma unroll
+        for (int p = 0; p < QK_ML8 / 2; p++) {
+            const uint8_t byte = blk->qs[p];
+            const uint8_t lo   = byte & 0x0F;
+            const uint8_t hi   = (byte >> 4) & 0x0F;
+            y_row[k_base + p * 2]     = ml8_fp8_e4m3_to_fp32(lut_g[lo]) * scale;
+            y_row[k_base + p * 2 + 1] = ml8_fp8_e4m3_to_fp32(lut_g[hi]) * scale;
+        }
+    }
+}
+
+void ggml_cuda_op_ml8_get_rows(
+    ggml_backend_cuda_context & ctx,
+    ggml_tensor *               dst) {
+    const ggml_tensor * w    = dst->src[0];
+    const ggml_tensor * cent = dst->src[1];
+    const ggml_tensor * ids  = dst->src[2];
+
+    GGML_ASSERT(w != nullptr && cent != nullptr && ids != nullptr);
+    GGML_ASSERT(w->type    == GGML_TYPE_ML8_4);
+    GGML_ASSERT(cent->type == GGML_TYPE_F8_E4M3);
+    GGML_ASSERT(ids->type  == GGML_TYPE_I32);
+    GGML_ASSERT(dst->type  == GGML_TYPE_F32);
+    GGML_ASSERT(ggml_is_contiguous(w));
+    GGML_ASSERT(ggml_is_contiguous(cent));
+    GGML_ASSERT(ggml_is_contiguous(ids));
+    GGML_ASSERT(ggml_is_contiguous(dst));
+
+    const int32_t K = (int32_t) w->ne[0];
+    const int32_t N = (int32_t) w->ne[1];
+    GGML_ASSERT(K % QK_ML8 == 0);
+    const int32_t n_groups_k = K / QK_ML8;
+    GGML_ASSERT(cent->ne[0] == 16);
+    GGML_ASSERT(cent->ne[1] == n_groups_k);
+    GGML_ASSERT(dst->ne[0] == K);
+
+    const int64_t nr = ggml_nelements(ids);
+    GGML_ASSERT(ggml_nrows(dst) == nr);
+    if (nr == 0) {
+        return;
+    }
+
+    cudaStream_t stream = ctx.stream();
+
+    const block_ml8_4 * w_d   = (const block_ml8_4 *) w->data;
+    const uint8_t     * lut_d = (const uint8_t     *) cent->data;
+    const int32_t     * ids_d = (const int32_t     *) ids->data;
+    float             * y_d   = (float             *) dst->data;
+
+    const int threads = (n_groups_k < 256) ? ((n_groups_k + 31) / 32) * 32 : 256;
+    const dim3 grid((unsigned) nr);
+    ml8_get_rows_kernel<<<grid, dim3(threads > 0 ? threads : 32), 0, stream>>>(
+        w_d, lut_d, ids_d, y_d, K, N, n_groups_k, nr);
+}
+
 // ─────────────────────────────────────────────────────────────────────
 // No-LUT FP8-WMMA mul_mat for scaled-fp8 weights (GGML_TYPE_ML8_FP8).
 //

ggml/src/ggml-cuda/ml8.cuh

@@ -122,6 +122,18 @@ void ggml_cuda_op_ml8_mul_mat(
     ggml_backend_cuda_context & ctx,
     ggml_tensor *               dst);
 
+// GGML_OP_ML8_GET_ROWS dispatch — native 4-bit token-embedding gather.
+// Unlike ggml_cuda_op_ml8_mul_mat this needs NO AITER GEMM: it gathers row
+// ids[i] from the ml8-4 weight and dequantizes via the per-K-group centroid
+// LUT directly on device. Available on any CUDA/HIP build.
+//   dst:        fp32  [K, ids->ne0, ids->ne1, ids->ne2]
+//   dst->src[0]: w    — GGML_TYPE_ML8_4,   ne[0]=K (mult. of QK_ML8=64), ne[1]=N(vocab)
+//   dst->src[1]: cent — GGML_TYPE_F8_E4M3, ne[0]=16, ne[1]=K/QK_ML8 (shared LUT)
+//   dst->src[2]: ids  — GGML_TYPE_I32
+void ggml_cuda_op_ml8_get_rows(
+    ggml_backend_cuda_context & ctx,
+    ggml_tensor *               dst);
+
 // Execute a plain GGML_OP_MUL_MAT whose src[0] is a GGML_TYPE_ML8_FP8
 // (scaled-fp8) weight, via the no-LUT FP8-WMMA path (WEIGHT_FORMAT=0).
 // Unlike GGML_OP_ML8_MUL_MAT there is no centroid sidecar:

ggml/src/ggml-ml8.c

@@ -327,3 +327,33 @@ struct ggml_tensor * ggml_ml8_mul_mat_id(
     y->src[3] = ids;
     return y;
 }
+
+struct ggml_tensor * ggml_ml8_get_rows(
+        struct ggml_context * ctx,
+        struct ggml_tensor  * w,
+        struct ggml_tensor  * centroids,
+        struct ggml_tensor  * ids) {
+    GGML_ASSERT(w         != NULL);
+    GGML_ASSERT(centroids != NULL);
+    GGML_ASSERT(ids       != NULL);
+    GGML_ASSERT(w->type         == GGML_TYPE_ML8_4);
+    GGML_ASSERT(centroids->type == GGML_TYPE_F8_E4M3);
+    GGML_ASSERT(ids->type       == GGML_TYPE_I32);
+
+    // w [K, N] (N = vocab rows), gather row ids[i] → K fp32 (one embedding vector).
+    const int64_t K = w->ne[0];
+    GGML_ASSERT(K % QK_ML8 == 0 && "K must be a multiple of QK_ML8=64");
+
+    const int64_t n_groups_k = K / QK_ML8;
+    GGML_ASSERT(centroids->ne[0] == 16);
+    GGML_ASSERT(centroids->ne[1] == n_groups_k);
+
+    // Output mirrors ggml_get_rows: [K, ids->ne[0], ids->ne[1], ids->ne[2]].
+    const int64_t ne[4] = { K, ids->ne[0], ids->ne[1], ids->ne[2] };
+    struct ggml_tensor * y = ggml_new_tensor(ctx, GGML_TYPE_F32, 4, ne);
+    y->op     = GGML_OP_ML8_GET_ROWS;
+    y->src[0] = w;
+    y->src[1] = centroids;
+    y->src[2] = ids;
+    return y;
+}

ggml/src/ggml.c

@@ -1179,9 +1179,10 @@ static const char * GGML_OP_NAME[GGML_OP_COUNT] = {
     "ML8_MUL_MAT",
     "ML8_APPLY_ROTATION",
     "ML8_MUL_MAT_ID",
+    "ML8_GET_ROWS",
 };
 
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
+static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103");
 
 static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "none",
@@ -1296,9 +1297,10 @@ static const char * GGML_OP_SYMBOL[GGML_OP_COUNT] = {
     "ml8_mul_mat(w,centroids,x)",
     "ml8_apply_rotation(x,h_a)",
     "ml8_mul_mat_id(w,centroids,x,ids)",
+    "ml8_get_rows(w,centroids,ids)",
 };
 
-static_assert(GGML_OP_COUNT == 102, "GGML_OP_COUNT != 102");
+static_assert(GGML_OP_COUNT == 103, "GGML_OP_COUNT != 103");
 
 static_assert(GGML_OP_POOL_COUNT == 2, "GGML_OP_POOL_COUNT != 2");