Merge f-q4-hipkg2 (da7b1d022): the AMD prompt-expert kernels (F16 K-quant, native-format fused IQ, Q2 blobs) follow the compiled geometry; default build unchanged (HIP gfx1201 182/182 and gfx1151 194/194 touched functions, CUDA SASS 3573/3573 and 160/160 host objects, SYCL 129/129 objects and 22/22 checksums, Metal 85/85 objects); HIP q4_lite gate 2 0.000183 vs floor 0.000173 (Qwen4 prep)

include/strata/prefill/moe_fused.hpp

@@ -15,6 +15,8 @@
 // rounding of the activations and of H) by about as much as MMQ differs from FP32 (tests/cuda/prefill_fused_moe_test).
 #pragma once
 
+#include "strata/core/geometry.hpp"
+
 #include <cstddef>
 #include <cstdint>
 
@@ -22,6 +24,11 @@ namespace strata::prefill::fused {
 
 /// Routed rows (token, k pairs of one expert) per tile.
 constexpr int kTileRows = 64;
+/// The compiled geometry's widths the expert kernels are written for: the hidden width (the K of gate/up, the rows of
+/// down), the expert width (the K of down, half the rows of gate/up) and the gate/up rows.  Flash-Next: 2560, 640, 1280.
+constexpr int kMoeEmbd = (int) strata::core::kCompiledGeometry.n_embd;
+constexpr int kMoeFF = (int) strata::core::kCompiledGeometry.n_ff;
+constexpr int kMoeGuRows = 2 * kMoeFF;
 /// Experts per launch: their blob pointers travel in the launch's parameters (1 KB).
 constexpr int kMaxBatch = 128;
 /// The room of WideBatch::blob (below): a whole layer of 512 experts, which STRATA_PF_BATCH=512 (opt-in) launches at
@@ -72,8 +79,8 @@ struct WideBatch {
     int e0 = 0, e1 = 0;
     const uint8_t* blob[kBatchRoom] = {};
 };
-/// The batch's products: gate/up from `xa` (quantize_act of the layer's input, [tokens][2560]) through the rows
-/// `src` of `group`, SwiGLU, H as int8 into `ha` (act_bytes(n, 640)), down into `dm` [n][2560] FP32 at the rows of
+/// The batch's products: gate/up from `xa` (quantize_act of the layer's input, [tokens][n_embd]) through the rows
+/// `src` of `group`, SwiGLU, H as int8 into `ha` (act_bytes(n, n_ff)), down into `dm` [n][n_embd] FP32 at the rows of
 /// `group`.  `n`: the layer's routed rows (tokens * k), which bounds the launch grid.
 void experts(const Batch& b, int n_expert, int64_t n, const void* scratch, const void* xa, const int32_t* src, void* ha,
              float* dm, void* stream);

include/strata/prefill/moe_fused_iq.hpp

@@ -19,7 +19,7 @@
 
 namespace strata::prefill::fused {
 
-/// One layer's native expert geometry (cpu::NativeFmt's fields the kernels read; n_embd 2560, n_ff 640).
+/// One layer's native expert geometry (cpu::NativeFmt's fields the kernels read; n_embd and n_ff are the compiled geometry's: 2560 and 640 for Flash-Next).
 struct NativeGeom {
     int gu_type = -1, d_type = -1;    ///< ggml types
     size_t gu_row = 0, d_row = 0;     ///< bytes per gate/up row and per down row
@@ -37,7 +37,7 @@ bool gfx12_iq2s_default();
 void quantize_act_native(const float* x, int64_t rows, int64_t cols, void* xa, void* stream);
 
 /// experts() for native blobs: b.blob[e - e0] is the native blob of expert e (gate at 0, up at g.up_off, down at
-/// g.down_off; 2-byte aligned).  `xa` from quantize_act_native; H in `ha` (act_bytes(n, 640)); down into `dm` at the
+/// g.down_off; 2-byte aligned).  `xa` from quantize_act_native; H in `ha` (act_bytes(n, n_ff)); down into `dm` at the
 /// rows of `group`.  `dm` 16-byte aligned (gfx12 stores float4; checked there - the prompt path's buffers are 256-aligned
 /// through Alloc's 256-byte takes from a 256-aligned cache slot or cudaMalloc).
 void experts_native(const Batch& b, const NativeGeom& g, int n_expert, int64_t n, const void* scratch, const void* xa,

src/prefill/kernels.cu

@@ -2270,28 +2270,31 @@ __global__ void route_kernel(const float* __restrict__ logits, int32_t* __restri
     selected_sum = fmaxf(warp_sum(selected_sum), 6.103515625e-5f);
     if (lane < KU) wout[t * KU + lane] = selected / selected_sum;
 }
-// Strata blob: gate/up codes [1280][640 B], down codes [2560][160 B], gate/up scales [1280][40] f16, down scales [2560][10] f16
+// Strata blob: gate/up codes [2 F][E / 4 B], down codes [E][F / 4 B], gate/up scales [2 F][E / 64] f16, down scales [E][F / 64] f16
+// (E = N, F = FFC: the compiled geometry's; Flash-Next 2560 and 640: [1280][640 B], [2560][160 B], [1280][40], [2560][10])
+constexpr int BLOB_GU_ROWS = 2 * (int) FFC, BLOB_ROW_GU = N / 4, BLOB_ROW_D = (int) FFC / 4, BLOB_SC_GU = N / 64, BLOB_SC_D = (int) FFC / 64;
+static_assert(N % 64 == 0 && FFC % 64 == 0, "the Strata Q2 blob holds whole 64-weight blocks");
 template <bool HALF>
 __global__ void blob_dequant_kernel(const uint8_t* __restrict__ blob, uint16_t* __restrict__ gu16,
                                     uint16_t* __restrict__ d16) {
-    constexpr size_t O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
-                     O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
+    constexpr size_t O_D_CODES = (size_t) BLOB_GU_ROWS * BLOB_ROW_GU, O_GU_SC = O_D_CODES + (size_t) N * BLOB_ROW_D,
+                     O_D_SC = O_GU_SC + (size_t) BLOB_GU_ROWS * BLOB_SC_GU * 2;
     const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;   // one thread per 4 weights (one code byte)
-    const int64_t n_gu = 1280LL * 640, n_d = 2560LL * 160;
+    const int64_t n_gu = (int64_t) BLOB_GU_ROWS * BLOB_ROW_GU, n_d = (int64_t) N * BLOB_ROW_D;
     if (i < n_gu) {
-        const int64_t row = i / 640, byte = i % 640;
-        const uint8_t c = blob[row * 640 + byte];
-        const uint8_t* sp = blob + O_GU_SC + (size_t) (row * 40 + (byte * 4) / 64) * 2;
+        const int64_t row = i / BLOB_ROW_GU, byte = i % BLOB_ROW_GU;
+        const uint8_t c = blob[row * BLOB_ROW_GU + byte];
+        const uint8_t* sp = blob + O_GU_SC + (size_t) (row * BLOB_SC_GU + (byte * 4) / 64) * 2;
         const float d = __half2float(__ushort_as_half((uint16_t) (sp[0] | (sp[1] << 8))));
-        uint16_t* o = gu16 + row * 2560 + byte * 4;
+        uint16_t* o = gu16 + row * N + byte * 4;
 #pragma unroll
         for (int k = 0; k < 4; ++k) { const float v = (float) (((c >> (2 * k)) & 3) - 1) * d; o[k] = HALF ? hf(v) : bf(v); }
     } else if (i < n_gu + n_d) {
-        const int64_t j = i - n_gu, row = j / 160, byte = j % 160;
-        const uint8_t c = blob[O_D_CODES + row * 160 + byte];
-        const uint8_t* sp = blob + O_D_SC + (size_t) (row * 10 + (byte * 4) / 64) * 2;
+        const int64_t j = i - n_gu, row = j / BLOB_ROW_D, byte = j % BLOB_ROW_D;
+        const uint8_t c = blob[O_D_CODES + row * BLOB_ROW_D + byte];
+        const uint8_t* sp = blob + O_D_SC + (size_t) (row * BLOB_SC_D + (byte * 4) / 64) * 2;
         const float d = __half2float(__ushort_as_half((uint16_t) (sp[0] | (sp[1] << 8))));
-        uint16_t* o = d16 + row * 640 + byte * 4;
+        uint16_t* o = d16 + row * FFC + byte * 4;
 #pragma unroll
         for (int k = 0; k < 4; ++k) { const float v = (float) (((c >> (2 * k)) & 3) - 1) * d; o[k] = HALF ? hf(v) : bf(v); }
     }
@@ -3212,11 +3215,11 @@ void route(const float* logits, int32_t* ids, float* weights, int64_t T, int64_t
     check("route");
 }
 void blob_dequant(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream) {
-    blob_dequant_kernel<false><<<blocks_for(1280LL * 640 + 2560LL * 160), 256, 0, (cudaStream_t) stream>>>(blob, gu16, down16);
+    blob_dequant_kernel<false><<<blocks_for((int64_t) BLOB_GU_ROWS * BLOB_ROW_GU + (int64_t) N * BLOB_ROW_D), 256, 0, (cudaStream_t) stream>>>(blob, gu16, down16);
     check("blob_dequant");
 }
 void blob_dequant_f16(const uint8_t* blob, uint16_t* gu16, uint16_t* down16, void* stream) {
-    blob_dequant_kernel<true><<<blocks_for(1280LL * 640 + 2560LL * 160), 256, 0, (cudaStream_t) stream>>>(blob, gu16, down16);
+    blob_dequant_kernel<true><<<blocks_for((int64_t) BLOB_GU_ROWS * BLOB_ROW_GU + (int64_t) N * BLOB_ROW_D), 256, 0, (cudaStream_t) stream>>>(blob, gu16, down16);
     check("blob_dequant_f16");
 }
 void swiglu_interleaved(const float* gu, uint16_t* h16, int64_t n, void* stream) {

src/prefill/moe_f16_common.cuh

@@ -47,8 +47,10 @@ namespace strata::prefill::fused::f16 {
 constexpr int AB = 80;           // activation / H bytes per 64 values: 64 codes, {d0, d1}, {sum0, sum1}
 constexpr int THREADS = 256;
 constexpr int ROWS = 128;        // weight rows per work item
-constexpr int GU_NFB = 640 / 64, D_NFB = 2560 / ROWS;   // work items per tile: 64 features, 128 output rows
-constexpr int GU_ACT_LD = 2560 / 64 * AB, H_LD = 640 / 64 * AB;   // bytes per activation / H row
+// The widths are the compiled geometry's (kMoeEmbd 2560 and kMoeFF 640 for Flash-Next: 10 gate/up items and 20 down items a tile).
+static_assert(kMoeFF % 64 == 0 && kMoeEmbd % ROWS == 0, "the F16 expert kernels: the expert width must be a multiple of 64 (a gate/up item is 64 features) and the hidden width of 128 (a down item is 128 output rows)");
+constexpr int GU_NFB = kMoeFF / 64, D_NFB = kMoeEmbd / ROWS;   // work items per tile: 64 features, 128 output rows
+constexpr int GU_ACT_LD = kMoeEmbd / 64 * AB, H_LD = kMoeFF / 64 * AB;   // bytes per activation / H row
 
 /// STRATA_PF_F16=1 (on top of STRATA_PF_FUSED=1): the covered formats take the F16 kernels on gfx11.
 inline bool requested() {
@@ -225,8 +227,8 @@ __device__ __forceinline__ void gu_epilogue(const v8f (&acc)[BN / 16], uint8_t*
 }
 
 /// Down epilogue (`lds` at least D_EPI_BYTES, free): each wave's 16 x 16 tile through LDS ([token][row]), then 16
-/// consecutive FP32 rows per routed row into dm [routed row][2560].
-/// D16 (STRATA_PF_DM16): the rows go out as FP16 (saturated at +-65504, 2,560 halves a row, two neighbouring columns in one 4-byte store from the even lane)
+/// consecutive FP32 rows per routed row into dm [routed row][kMoeEmbd].
+/// D16 (STRATA_PF_DM16): the rows go out as FP16 (saturated at +-65504, kMoeEmbd halves a row, two neighbouring columns in one 4-byte store from the even lane)
 /// for moe_combine16 - half the bytes written here and read there.
 template <int BN, bool D16 = false>
 __device__ __forceinline__ void down_epilogue(const v8f (&acc)[BN / 16], uint8_t* lds, int tid, int live_tiles, int nrows, int row0,
@@ -247,10 +249,10 @@ __device__ __forceinline__ void down_epilogue(const v8f (&acc)[BN / 16], uint8_t
                 const float v = fminf(fmaxf(tile[flat], -65504.0f), 65504.0f);
                 const float nx = __shfl_down(v, 1);
                 if (t < nrows && !(flat & 1))
-                    reinterpret_cast<uint32_t*>(reinterpret_cast<uint16_t*>(dm) + (size_t) (row0 + t) * 2560 + r0 + (flat & 15))[0] =
+                    reinterpret_cast<uint32_t*>(reinterpret_cast<uint16_t*>(dm) + (size_t) (row0 + t) * kMoeEmbd + r0 + (flat & 15))[0] =
                         (uint32_t) __half_as_ushort(__float2half(v)) | ((uint32_t) __half_as_ushort(__float2half(nx)) << 16);
             } else {
-                if (t < nrows) dm[(size_t) (row0 + t) * 2560 + r0 + (flat & 15)] = tile[flat];
+                if (t < nrows) dm[(size_t) (row0 + t) * kMoeEmbd + r0 + (flat & 15)] = tile[flat];
             }
         }
         __builtin_amdgcn_wave_barrier();

src/prefill/moe_f16_kq.cu

@@ -41,7 +41,7 @@ namespace {
 using namespace f16;
 
 constexpr int T_Q4_K = GGML_TYPE_Q4_K, T_Q5_K = GGML_TYPE_Q5_K;
-constexpr int GU_NS = 2560 / 64, D_NS = 640 / 64;   // stages of 64 values
+constexpr int GU_NS = kMoeEmbd / 64, D_NS = kMoeFF / 64;   // stages of 64 values (the compiled geometry's: 40 and 10 for Flash-Next)
 constexpr int FEAT = 64;                             // features (gate rows = up rows) per gate/up item
 
 // Byte i of the 16-byte K-quant block header (d, dmin, scales[12]).
@@ -213,7 +213,7 @@ __global__ void __launch_bounds__(THREADS) gate_up_kernel(const Batch b, const N
 #endif
 }
 
-// Q5_1 down: 128 of the 2560 output rows against a tile's (or pair's) H rows, K = 640 in 10 stages; the routed rows'
+// Q5_1 down: 128 of the kMoeEmbd output rows against a tile's (or pair's) H rows, K = kMoeFF in kMoeFF / 64 stages; the routed rows'
 // FP32 outputs in `dm` (native_w11_kernel<T_Q5_1, false>'s).
 template <int BN, bool D16 = false>
 __global__ void __launch_bounds__(THREADS) down_kernel(const Batch b, const NativeGeom geo, const int32_t* __restrict__ ts,

src/prefill/moe_fused.cu

@@ -44,17 +44,23 @@ void ck(cudaError_t e, const char* what) {
 constexpr int AB = 80;                  // activation bytes per 64 values: 64 codes, then {d0, c0, d1, c1}
 constexpr int MAGIC = 0x4B400000;       // the bits of 1.5 * 2^23
 constexpr float MAGICF = 12582912.0f;
-// The Strata Q2_0 blob (moe_mmq.cu strata_q2_kernel, kernels.cu blob_dequant_kernel): gate/up codes [1280][640 B]
-// (rows interleaved: gate of feature f at row 2f, its up at 2f + 1), down codes [2560][160 B], gate/up scales
-// [1280][40] fp16, down scales [2560][10] fp16.
-constexpr size_t O_GU_CODES = 0, O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
-                 O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
+// The Strata Q2_0 blob (moe_mmq.cu strata_q2_kernel, kernels.cu blob_dequant_kernel): gate/up codes [2 F][E / 4 B]
+// (rows interleaved: gate of feature f at row 2f, its up at 2f + 1), down codes [E][F / 4 B], gate/up scales
+// [2 F][E / 64] fp16, down scales [E][F / 64] fp16; E, F the compiled geometry's (Flash-Next 2560, 640: [1280][640 B],
+// [2560][160 B], [1280][40], [2560][10]).
+constexpr int GU_ROWS = kMoeGuRows, D_ROWS = kMoeEmbd;           // weight rows of the two products
+constexpr int ROW_GU = kMoeEmbd / 4, ROW_D = kMoeFF / 4;         // code bytes per weight row
+constexpr int SC_GU = kMoeEmbd / 64, SC_D = kMoeFF / 64;         // 64-weight blocks (scales) per row = stages along K
+constexpr int H_ROW_BYTES = SC_D * AB;                           // bytes of a row of H (AB per 64 values)
+static_assert(kMoeEmbd % 64 == 0 && kMoeFF % 64 == 0, "the Strata Q2 blob holds whole 64-weight blocks");
+constexpr size_t O_GU_CODES = 0, O_D_CODES = (size_t) GU_ROWS * ROW_GU, O_GU_SC = O_D_CODES + (size_t) D_ROWS * ROW_D,
+                 O_D_SC = O_GU_SC + (size_t) GU_ROWS * SC_GU * 2;
 constexpr int THREADS = 512;            // 16 warps: 4 along the weight rows (64 each) x 4 along the tile rows (16 each)
 constexpr int WROWS = 256;              // weight rows per work item: 128 features of gate/up, or 256 outputs of down
 constexpr int STAGES = 4;               // cp.async pipeline depth, one 64-weight block of K per stage
 constexpr int STAGE_BYTES = WROWS * 16 + kTileRows * AB;
 constexpr size_t smem_bytes(bool gu) {
-    return (size_t) WROWS * (gu ? 40 : 10) * 2 + STAGES * STAGE_BYTES + kTileRows * 4;
+    return (size_t) WROWS * (gu ? SC_GU : SC_D) * 2 + STAGES * STAGE_BYTES + kTileRows * 4;
 }
 
 // the stored position of value k of a 32-value half-block (see "The K order")
@@ -220,7 +226,7 @@ __device__ __forceinline__ void mma_s8(int (&d)[4], const uint32_t (&a)[4], uint
 
 // One work item = (tile of up to 64 routed rows of expert e, a block of 256 weight rows); a persistent grid walks
 // the batch's items (their count is on the device).  GU: gate/up rows of 128 features against the rows' tokens'
-// activations (`act` per token, 2560 values), SwiGLU, H to int8 per 32 features into `out` (per row, 640 values).
+// activations (`act` per token, E values), SwiGLU, H to int8 per 32 features into `out` (per row, F values).
 // !GU: down rows (256 outputs) against H (`act` per row), FP32 into `dm` per row.
 // Warp (wf, wt): weight rows 64 wf.. as four m16 tiles (GU: a tile = 8 features, its gate rows as mma rows 0-7 and
 // its up rows as 8-15, so a lane holds gate and up of one feature: SwiGLU in registers), tile rows 16 wt.. as two n8.
@@ -229,11 +235,14 @@ __global__ void __launch_bounds__(THREADS, 1)
 expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, const int32_t* __restrict__ src,
               uint8_t* __restrict__ out, float* __restrict__ dm) {
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 750
-    constexpr int NKB = GU ? 40 : 10;                 // 64-weight blocks along K
-    constexpr int NFB = GU ? 5 : 10;                  // weight-row blocks per tile
-    constexpr int CODE_LD = GU ? 640 : 160;           // code bytes per weight row
+    constexpr int NKB = GU ? SC_GU : SC_D;            // 64-weight blocks along K
+    constexpr int NFB = (GU ? GU_ROWS : D_ROWS) / WROWS;   // weight-row blocks per tile
+    static_assert(GU_ROWS % WROWS == 0 && D_ROWS % WROWS == 0, "the Strata Q2 expert kernel: weight rows are whole work items of 256");
+    static_assert((SC_GU * 2) % 16 == 0, "the Strata Q2 expert kernel stages a row's gate/up scales in 16-byte copies (hidden width a multiple of 512)");
+    constexpr int NSC = GU ? SC_GU * 2 / 16 : SC_D * 2 / 4;   // scale copies per weight row (16 B for gate/up, 4 B for down)
+    constexpr int CODE_LD = GU ? ROW_GU : ROW_D;      // code bytes per weight row
     constexpr size_t O_CODES = GU ? O_GU_CODES : O_D_CODES, O_SC = GU ? O_GU_SC : O_D_SC;
-    constexpr int ACT_LD = NKB * AB;                  // bytes per activation row: 3200 (a token), 800 (a row's H)
+    constexpr int ACT_LD = NKB * AB;                  // bytes per activation row: 3200 (a token), 800 (a row's H) for Flash-Next
     constexpr unsigned M2 = 0x03030303u;
     extern __shared__ __align__(16) uint8_t smem[];
     const __half* wsc = (const __half*) smem;                     // [256][NKB] the work item's weight scales
@@ -259,10 +268,10 @@ expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, c
             const int r = row0 + min(tid, nrows - 1);             // rows past the tile's end repeat its last one
             srow[tid] = GU ? src[r] : r;
         }
-        for (int c = tid; c < WROWS * 5; c += THREADS) {          // the scales: 5 x 16 B (GU) or 5 x 4 B a row
-            const int r = c / 5, q = c % 5;
-            if (GU) cp16(smem + r * 80 + q * 16, blob + O_SC + (size_t) (rbase + r) * 80 + q * 16);
-            else cp4(smem + r * 20 + q * 4, blob + O_SC + (size_t) (rbase + r) * 20 + q * 4);
+        for (int c = tid; c < WROWS * NSC; c += THREADS) {        // the scales: NSC x 16 B (GU) or NSC x 4 B a row (5 for Flash-Next)
+            const int r = c / NSC, q = c % NSC;
+            if (GU) cp16(smem + r * (SC_GU * 2) + q * 16, blob + O_SC + (size_t) (rbase + r) * (SC_GU * 2) + q * 16);
+            else cp4(smem + r * (SC_D * 2) + q * 4, blob + O_SC + (size_t) (rbase + r) * (SC_D * 2) + q * 4);
         }
         __syncthreads();                                          // srow
         auto load = [&](int kb) {
@@ -356,7 +365,7 @@ expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, c
                     for (int o = 4; o < 32; o <<= 1) sum += __shfl_xor_sync(0xffffffffu, sum, o);
                     const int r = nb0 + 8 * n + 2 * tig + c;
                     if (r < nrows) {
-                        uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (blk >> 1) * AB;
+                        uint8_t* o = out + (size_t) (row0 + r) * H_ROW_BYTES + (blk >> 1) * AB;
                         const int hh = blk & 1;
 #pragma unroll
                         for (int i = 0; i < 4; ++i) o[32 * hh + perm32(8 * i + g)] = (uint8_t) (int8_t) qv[i];
@@ -373,7 +382,7 @@ expert_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, c
 #pragma unroll
                     for (int q = 0; q < 4; ++q) {
                         const int r = nb0 + 8 * n + 2 * tig + (q & 1);
-                        if (r < nrows) dm[(size_t) (row0 + r) * 2560 + rbase + (q < 2 ? ra[i] : rb[i])] = acc[i][n][q];
+                        if (r < nrows) dm[(size_t) (row0 + r) * kMoeEmbd + rbase + (q < 2 ? ra[i] : rb[i])] = acc[i][n][q];
                     }
                 }
         }
@@ -417,9 +426,10 @@ __global__ void __launch_bounds__(W_THREADS)
 expert_w11_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ act, const int32_t* __restrict__ src,
                   uint8_t* __restrict__ out, float* __restrict__ dm) {
 #if STRATA_FUSED_W11
-    constexpr int NKB = GU ? 40 : 10;                 // 64-weight blocks along K
-    constexpr int NFB = (GU ? 1280 : 2560) / W_ROWS;  // work items per tile
-    constexpr int CODE_LD = GU ? 640 : 160;           // code bytes per weight row
+    constexpr int NKB = GU ? SC_GU : SC_D;            // 64-weight blocks along K
+    constexpr int NFB = (GU ? GU_ROWS : D_ROWS) / W_ROWS;  // work items per tile
+    static_assert(GU_ROWS % W_ROWS == 0 && D_ROWS % W_ROWS == 0, "the Strata Q2 gfx11 kernel: weight rows are whole work items of 128");
+    constexpr int CODE_LD = GU ? ROW_GU : ROW_D;      // code bytes per weight row
     constexpr size_t O_CODES = GU ? O_GU_CODES : O_D_CODES, O_SC = GU ? O_GU_SC : O_D_SC;
     constexpr int ACT_LD = NKB * AB;                  // bytes per activation row
     constexpr unsigned M2 = 0x03030303u;
@@ -555,7 +565,7 @@ expert_w11_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ ac
 #pragma unroll 8
                     for (int j = 0; j < 32; ++j) am = fmaxf(am, fabsf(hs[r][32 * hh + j]));
                     const float inv = am > 0.0f ? 127.0f / am : 0.0f;
-                    uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (size_t) fb * AB;
+                    uint8_t* o = out + (size_t) (row0 + r) * H_ROW_BYTES + (size_t) fb * AB;
                     int sum = 0;
                     uint32_t wd[8] = {};
 #pragma unroll
@@ -578,7 +588,7 @@ expert_w11_kernel(const Batch b, const Tables tb, const uint8_t* __restrict__ ac
                 for (int nt = 0; nt < 2; ++nt) {
                     const int r = 32 * wn + 16 * nt + l16;
                     if (r < nrows) {
-                        float* d = dm + (size_t) (row0 + r) * 2560 + rbase + 32 * wm + 16 * mt + hi;
+                        float* d = dm + (size_t) (row0 + r) * kMoeEmbd + rbase + 32 * wm + 16 * mt + hi;
 #pragma unroll
                         for (int i = 0; i < 8; ++i) d[2 * i] = acc[mt][nt][i];
                     }
@@ -735,13 +745,13 @@ void experts(const Batch& b, int n_expert, int64_t n, const void* scratch, const
     // the most tiles the batch can have: every row in it, plus a partial tile per expert
     const int64_t tiles = (n + kTileRows - 1) / kTileRows + (b.e1 - b.e0);
 #if defined(__HIPCC__)
-    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (1280 / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_gu));
-    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (2560 / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_d));
+    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (GU_ROWS / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_gu));
+    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (D_ROWS / W_ROWS), (int64_t) d.sms * wgp_blocks(d.occ_d));
     expert_w11_kernel<true><<<g_gu, W_THREADS, 0, s>>>(b, tb, (const uint8_t*) xa, src, (uint8_t*) ha, nullptr);
     expert_w11_kernel<false><<<g_d, W_THREADS, 0, s>>>(b, tb, (const uint8_t*) ha, src, nullptr, dm);
 #else
-    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * 5, (int64_t) d.sms * d.occ_gu);
-    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * 10, (int64_t) d.sms * d.occ_d);
+    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (GU_ROWS / WROWS), (int64_t) d.sms * d.occ_gu);
+    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (D_ROWS / WROWS), (int64_t) d.sms * d.occ_d);
     expert_kernel<true><<<g_gu, THREADS, smem_bytes(true), s>>>(b, tb, (const uint8_t*) xa, src, (uint8_t*) ha,
                                                                 nullptr);
     expert_kernel<false><<<g_d, THREADS, smem_bytes(false), s>>>(b, tb, (const uint8_t*) ha, src, nullptr, dm);

src/prefill/moe_fused_iq.cu

@@ -76,7 +76,11 @@ constexpr bool T75_DEV = true;    // the Q2_0 stage epilogue in native_kernel: f
 constexpr int ASTAGES_DEV = ASTAGES;
 constexpr bool T75_DEV = false;
 #endif
-constexpr int GU_ROWS_K = 2560, D_ROWS_K = 640;   // K of gate/up (n_embd) and of down (n_ff)
+constexpr int GU_ROWS_K = kMoeEmbd, D_ROWS_K = kMoeFF;   // K of gate/up (n_embd) and of down (n_ff): the compiled geometry's (2560 and 640 for Flash-Next)
+// the weight rows of a product: gate/up 2 n_ff (1280), down n_embd (2560)
+constexpr int GU_ROWS_N = kMoeGuRows, D_ROWS_N = kMoeEmbd;
+constexpr int H_ROW_BYTES = D_ROWS_K / 64 * AB;   // bytes of a row of H in the int8 activation layout (800 for Flash-Next)
+static_assert(GU_ROWS_K % 256 == 0 && D_ROWS_K % 64 == 0, "the native expert kernels: K of gate/up is whole 256-value super-blocks, K of down whole 64-value stages");
 
 // the formats (ggml type ids)
 constexpr int T_IQ2_XXS = GGML_TYPE_IQ2_XXS, T_IQ2_XS = GGML_TYPE_IQ2_XS, T_IQ2_S = GGML_TYPE_IQ2_S,
@@ -587,7 +591,8 @@ native_kernel(const BatchOf<NT> b, const NativeGeom geo, const Tables tb, const
     constexpr int WT_BYTES = WR * WLDX, WS_FLOATS = F16 ? 0 : WR * 4, ACT_STAGE = TR * AB;
     constexpr bool TWO_PASS = NT != THREADS;   // the two-pass loads (load_unit_l) only in the opt-in NT = 256 kernels
     constexpr int NS = (GU ? GU_ROWS_K : D_ROWS_K) / 64;   // 64-value stages along K
-    constexpr int NFB = (GU ? 1280 : 2560) / WR;          // weight-row blocks per tile
+    constexpr int NFB = (GU ? GU_ROWS_N : D_ROWS_N) / WR;          // weight-row blocks per tile
+    static_assert(GU_ROWS_N % WR == 0 && D_ROWS_N % WR == 0, "the weight rows of a product are whole work items (2 n_ff and n_embd are multiples of 64 WW)");
     constexpr int ACT_LD = NS * AB;                       // bytes per activation row: 3200 (a token), 800 (a row's H)
     constexpr int BS = block_bytes(WT);
     constexpr bool K16 = per16(WT);
@@ -849,7 +854,7 @@ native_kernel(const BatchOf<NT> b, const NativeGeom geo, const Tables tb, const
                     const float inv = am > 0.0f ? 127.0f / am : 0.0f;
                     const int r = nb0 + 8 * n + 2 * tig + c;
                     if (r < nrows) {
-                        uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (blk >> 1) * AB;
+                        uint8_t* o = out + (size_t) (row0 + r) * H_ROW_BYTES + (blk >> 1) * AB;
                         const int hh = blk & 1;
 #pragma unroll
                         for (int i = 0; i < 4; ++i) o[32 * hh + 8 * i + g] = (uint8_t) (int8_t) __float2int_rn(h[i] * inv);
@@ -866,7 +871,7 @@ native_kernel(const BatchOf<NT> b, const NativeGeom geo, const Tables tb, const
 #pragma unroll
                     for (int q = 0; q < 4; ++q) {
                         const int r = nb0 + 8 * n + 2 * tig + (q & 1);
-                        if (r < nrows) dm[(size_t) (row0 + r) * 2560 + rbase + 64 * wf + 16 * i + 8 * (q >> 1) + g] = acc[i][n][q];
+                        if (r < nrows) dm[(size_t) (row0 + r) * kMoeEmbd + rbase + 64 * wf + 16 * i + 8 * (q >> 1) + g] = acc[i][n][q];
                     }
                 }
         }
@@ -927,7 +932,7 @@ native_ws_kernel(const WideBatch b, const NativeGeom geo, const Tables tb, const
     constexpr int WS_UPT = 256 / WS_PROD;   // decode units (sub-blocks) a producer thread takes per stage
     constexpr int TR = 64, WR = 128, WW = 2;
     constexpr int NS = (GU ? GU_ROWS_K : D_ROWS_K) / 64;
-    constexpr int NFB = (GU ? 1280 : 2560) / WR;
+    constexpr int NFB = (GU ? GU_ROWS_N : D_ROWS_N) / WR;
     constexpr int ACT_LD = NS * AB;
     constexpr int BS = block_bytes(WT);
     constexpr int WT_BYTES = WR * WLDH, ACT_STAGE = TR * AB;
@@ -1177,7 +1182,7 @@ native_ws_kernel(const WideBatch b, const NativeGeom geo, const Tables tb, const
                     const float inv = am > 0.0f ? 127.0f / am : 0.0f;
                     const int r = nb0 + 8 * n + 2 * tig + c;
                     if (r < nrows) {
-                        uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (blk >> 1) * AB;
+                        uint8_t* o = out + (size_t) (row0 + r) * H_ROW_BYTES + (blk >> 1) * AB;
                         const int hh = blk & 1;
 #pragma unroll
                         for (int i = 0; i < 4; ++i) o[32 * hh + 8 * i + g] = (uint8_t) (int8_t) __float2int_rn(h[i] * inv);
@@ -1195,7 +1200,7 @@ native_ws_kernel(const WideBatch b, const NativeGeom geo, const Tables tb, const
 #pragma unroll
                     for (int c = 0; c < 4; ++c) {
                         const int r = nb0 + 8 * n + 2 * tig + (c & 1);
-                        if (r < nrows) dm[(size_t) (row0 + r) * 2560 + rbase + 64 * wf + 16 * i + 8 * (c >> 1) + g] = acc[i][n][c];
+                        if (r < nrows) dm[(size_t) (row0 + r) * kMoeEmbd + rbase + 64 * wf + 16 * i + 8 * (c >> 1) + g] = acc[i][n][c];
                     }
                 }
         }
@@ -1313,7 +1318,7 @@ native_w11_kernel(const Batch b, const NativeGeom geo, const Tables tb, const ui
                   const int32_t* __restrict__ src, uint8_t* __restrict__ out, float* __restrict__ dm) {
 #if STRATA_NAT_WMMA
     constexpr int NS = (GU ? GU_ROWS_K : D_ROWS_K) / 64;  // 64-value stages along K
-    constexpr int NFB = (GU ? 1280 : 2560) / NW_ROWS;     // work items per tile
+    constexpr int NFB = (GU ? GU_ROWS_N : D_ROWS_N) / NW_ROWS;     // work items per tile
     constexpr int ACT_LD = NS * AB;
     constexpr int BS = block_bytes(WT);
     constexpr bool K16 = per16(WT);
@@ -1511,7 +1516,7 @@ native_w11_kernel(const Batch b, const NativeGeom geo, const Tables tb, const ui
 #pragma unroll 8
                     for (int j = 0; j < 32; ++j) am = fmaxf(am, fabsf(hs[r][32 * hh + j]));
                     const float inv = am > 0.0f ? 127.0f / am : 0.0f;
-                    uint8_t* o = out + (size_t) (row0 + r) * (10 * AB) + (size_t) fb * AB;
+                    uint8_t* o = out + (size_t) (row0 + r) * H_ROW_BYTES + (size_t) fb * AB;
                     uint32_t wd[8] = {};
                     int csum = 0;
 #pragma unroll
@@ -1537,11 +1542,11 @@ native_w11_kernel(const Batch b, const NativeGeom geo, const Tables tb, const ui
                         if (r < nrows) {
 #if STRATA_NAT_W12
                             // columns rbase + 32 wm + 16 mt + 8 hi .. + 7: contiguous, 32-byte aligned (dm: 16, checked in experts_native)
-                            float4* d = reinterpret_cast<float4*>(dm + (size_t) (row0 + r) * 2560 + rbase + 32 * wm + 16 * mt + 8 * hi);
+                            float4* d = reinterpret_cast<float4*>(dm + (size_t) (row0 + r) * kMoeEmbd + rbase + 32 * wm + 16 * mt + 8 * hi);
                             d[0] = make_float4(acc[mt][nt][0], acc[mt][nt][1], acc[mt][nt][2], acc[mt][nt][3]);
                             d[1] = make_float4(acc[mt][nt][4], acc[mt][nt][5], acc[mt][nt][6], acc[mt][nt][7]);
 #else
-                            float* d = dm + (size_t) (row0 + r) * 2560 + rbase + 32 * wm + 16 * mt + hi;
+                            float* d = dm + (size_t) (row0 + r) * kMoeEmbd + rbase + 32 * wm + 16 * mt + hi;
 #pragma unroll
                             for (int i = 0; i < 8; ++i) d[2 * i] = acc[mt][nt][i];
 #endif
@@ -2127,9 +2132,9 @@ void experts_native(const Batch& b, const NativeGeom& g, int n_expert, int64_t n
             std::fprintf(stderr, "prefill fused experts (native): the down rows (%p) must be 16-byte aligned\n", (void*) dm);
             std::exit(1);
         }
-        const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (1280 / NW_ROWS),
+        const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (GU_ROWS_N / NW_ROWS),
                                                            (int64_t) d.sms * (occ_gu ? occ_blocks(d.w12, g.gu_type, true) : wgp_blocks(d.occ)));
-        const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (2560 / NW_ROWS),
+        const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (D_ROWS_N / NW_ROWS),
                                                           (int64_t) d.sms * (occ_d ? occ_blocks(d.w12, g.d_type, false) : wgp_blocks(d.occ)));
         const uint8_t* xa8 = (const uint8_t*) xa;
         uint8_t* ha8 = (uint8_t*) ha;
@@ -2206,8 +2211,8 @@ void experts_native(const Batch& b, const NativeGeom& g, int n_expert, int64_t n
     // the most 64-row tiles the batch can have (every row in it, plus a partial tile per expert) - the items of the
     // ones inside a larger item end at once
     const int64_t tiles = (n + kTileRows - 1) / kTileRows + (b.e1 - b.e0);
-    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (1280 / weight_rows(ww)), (int64_t) d.sms * d.occ);
-    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (2560 / weight_rows(ww)), (int64_t) d.sms * d.occ);
+    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (GU_ROWS_N / weight_rows(ww)), (int64_t) d.sms * d.occ);
+    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (D_ROWS_N / weight_rows(ww)), (int64_t) d.sms * d.occ);
     switch (g.gu_type) {
         case T_IQ2_XXS: launch<T_IQ2_XXS, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
         case T_IQ2_XS: launch<T_IQ2_XS, true>(ww, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
@@ -2240,8 +2245,8 @@ void experts_native(const WideBatch& b, const NativeGeom& g, int n_expert, int64
     const int64_t tiles = (n + kTileRows - 1) / kTileRows + (b.e1 - b.e0);
     const bool ws = d.ws;
     const int occ = ws ? 1 : d.occ_nt;
-    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (1280 / weight_rows(2)), (int64_t) d.sms * occ);
-    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (2560 / weight_rows(2)), (int64_t) d.sms * occ);
+    const unsigned g_gu = (unsigned) std::min<int64_t>(tiles * (GU_ROWS_N / weight_rows(2)), (int64_t) d.sms * occ);
+    const unsigned g_d = (unsigned) std::min<int64_t>(tiles * (D_ROWS_N / weight_rows(2)), (int64_t) d.sms * occ);
     switch (g.gu_type) {
         case T_IQ2_XXS: launch<T_IQ2_XXS, true>(ws, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;
         case T_IQ2_XS: launch<T_IQ2_XS, true>(ws, g_gu, b, g, tb, xa, src, ha, nullptr, s); break;

src/prefill/moe_mmq.cu

@@ -9,6 +9,7 @@
 #if defined(STRATA_MMQ_TAB) && !defined(__HIPCC__)
 #include "moe_mmq_tab.hpp"   // declarations only: the tab kernels live in the per-type mmq-tab-*.cu TUs
 #endif
+#include "strata/prefill/moe_fused.hpp"   // kMoeEmbd, kMoeFF: the compiled geometry's widths
 #include "strata/prefill/moe_iq4_w4a8.hpp"
 
 #include "strata/prefill/mmq_kq_abi.hpp"
@@ -108,26 +109,29 @@ __global__ void copy1_kernel(const uint8_t* __restrict__ a, int64_t na, const ui
     else if (i < na + nb + nc) c_dst[i - na - nb] = c[i - na - nb];
 }
 
-// Strata blob: gate/up codes [1280][640 B], down codes [2560][160 B], gate/up scales [1280][40] f16, down scales
-// [2560][10] f16 (the layout of prefill/kernels.cu's blob_dequant_kernel).  A GGUF Q2_0 block is {f16 d; 16 code
+// Strata blob: gate/up codes [2 F][E / 4 B], down codes [E][F / 4 B], gate/up scales [2 F][E / 64] f16, down scales
+// [E][F / 64] f16 (the layout of prefill/kernels.cu's blob_dequant_kernel; E, F the compiled geometry's: Flash-Next
+// [1280][640 B], [2560][160 B], [1280][40], [2560][10]).  A GGUF Q2_0 block is {f16 d; 16 code
 // bytes} with the same 2-bit codes in the same order, so a block is a scale and a 16-byte run of codes.
+constexpr int Q2_GU_ROWS = fused::kMoeGuRows, Q2_ROW_GU = fused::kMoeEmbd / 4, Q2_ROW_D = fused::kMoeFF / 4, Q2_SC_GU = fused::kMoeEmbd / 64, Q2_SC_D = fused::kMoeFF / 64;
+static_assert(fused::kMoeEmbd % 64 == 0 && fused::kMoeFF % 64 == 0, "the Strata Q2 blob holds whole 64-weight blocks");
 __global__ void strata_q2_kernel(const uint8_t* __restrict__ blob, uint16_t* __restrict__ gu, uint16_t* __restrict__ dn) {
-    constexpr size_t O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
-                     O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
+    constexpr size_t O_D_CODES = (size_t) Q2_GU_ROWS * Q2_ROW_GU, O_GU_SC = O_D_CODES + (size_t) fused::kMoeEmbd * Q2_ROW_D,
+                     O_D_SC = O_GU_SC + (size_t) Q2_GU_ROWS * Q2_SC_GU * 2;
     const int64_t i = (int64_t) blockIdx.x * blockDim.x + threadIdx.x;   // one GGUF block
-    const int64_t n_gu = 1280LL * 40, n_d = 2560LL * 10;
+    const int64_t n_gu = (int64_t) Q2_GU_ROWS * Q2_SC_GU, n_d = (int64_t) fused::kMoeEmbd * Q2_SC_D;
     const uint8_t* codes;
     const uint16_t* scale;
     uint16_t* out;
     if (i < n_gu) {
-        const int64_t row = i / 40, b = i % 40;
-        codes = blob + row * 640 + b * 16;
-        scale = (const uint16_t*) (blob + O_GU_SC) + row * 40 + b;
+        const int64_t row = i / Q2_SC_GU, b = i % Q2_SC_GU;
+        codes = blob + row * Q2_ROW_GU + b * 16;
+        scale = (const uint16_t*) (blob + O_GU_SC) + row * Q2_SC_GU + b;
         out = gu + i * 9;
     } else if (i < n_gu + n_d) {
-        const int64_t j = i - n_gu, row = j / 10, b = j % 10;
-        codes = blob + O_D_CODES + row * 160 + b * 16;
-        scale = (const uint16_t*) (blob + O_D_SC) + row * 10 + b;
+        const int64_t j = i - n_gu, row = j / Q2_SC_D, b = j % Q2_SC_D;
+        codes = blob + O_D_CODES + row * Q2_ROW_D + b * 16;
+        scale = (const uint16_t*) (blob + O_D_SC) + row * Q2_SC_D + b;
         out = dn + j * 9;
     } else {
         return;
@@ -742,7 +746,7 @@ bool gather_native_group(const GatherGroup& g, size_t up_off, size_t gu_half_byt
 }
 
 void gather_strata_q2(const uint8_t* blob, void* gu_dst, void* d_dst, void* stream) {
-    strata_q2_kernel<<<blocks(1280LL * 40 + 2560LL * 10), 256, 0, (cudaStream_t) stream>>>(blob, (uint16_t*) gu_dst,
+    strata_q2_kernel<<<blocks((int64_t) Q2_GU_ROWS * Q2_SC_GU + (int64_t) fused::kMoeEmbd * Q2_SC_D), 256, 0, (cudaStream_t) stream>>>(blob, (uint16_t*) gu_dst,
                                                                                          (uint16_t*) d_dst);
     ck(cudaGetLastError(), "gather_strata_q2");
 }

tests/cuda/prefill_fused_iq_test.cpp

@@ -13,7 +13,7 @@
 // different expert ranges.  Part 2 (timing): one layer at the engine's chunks (2048, 3584, 8192 tokens; 512 experts,
 // top 10) on both paths and on MMQ's products alone (no gathers), for the IQ2_XS, IQ3_XXS and IQ3_S packs' most
 // common layers, with the grouping checked and a sample of rows against the reference.  --no-ref, --no-timing,
-// --only=NAME, --chunks=A,B, --unset-pf-fused (STRATA_PF_FUSED left unset instead of set to 1: the engine's default, for the
+// --only=NAME, --edge, --chunks=A,B, --unset-pf-fused (STRATA_PF_FUSED left unset instead of set to 1: the engine's default, for the
 // STRATA_PF_FUSED_Q2N=1 ctest entry).  Exit 77 without a CUDA device of sm_80 or newer.
 #include "strata/core/emulate.hpp"
 #include "strata/prefill/moe_fused_iq.hpp"
@@ -40,7 +40,8 @@ namespace {
 namespace mmq = strata::prefill::mmq;
 namespace fused = strata::prefill::fused;
 
-constexpr int N = 2560, FF = 640, K = 10, GROUP = 16;
+// The compiled geometry's widths (Flash-Next: 2560 x 640, 512 experts, top 10): build with -DSTRATA_GEOMETRY=<name> to test another shape.
+constexpr int N = fused::kMoeEmbd, FF = fused::kMoeFF, FF2 = fused::kMoeGuRows, K = (int) strata::core::kCompiledGeometry.n_expert_used, GROUP = 16;
 
 void ck(cudaError_t e, const char* what) {
     if (e != cudaSuccess) throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(e));
@@ -114,7 +115,7 @@ std::vector<uint8_t> make_blob(const Pair& p, const Geo& g, std::mt19937& rng) {
     return b;
 }
 
-// ggml's dequantized weights: gate [640][2560], up [640][2560], down [2560][640]
+// ggml's dequantized weights: gate [FF][N], up [FF][N], down [N][FF]
 void dequant(const Pair& p, const Geo& g, const std::vector<uint8_t>& b, std::vector<float>& gu, std::vector<float>& dn) {
     const auto* tg = ggml_get_type_traits(p.gu);
     const auto* td = ggml_get_type_traits(p.d);
@@ -178,9 +179,9 @@ Routing sort_rows(const std::vector<int32_t>& ids, int E) {
 struct MmqBufs {
     Dev xq, gu, h, hq, dm, ident, bounds, grp_gu, grp_d;
     MmqBufs(const Pair& p, int64_t rows, int E)
-        : xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * 1280 * 4), h((size_t) rows * FF * 4),
+        : xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * FF2 * 4), h((size_t) rows * FF * 4),
           hq(mmq::q8_bytes(rows, FF)), dm((size_t) rows * N * 4), ident((size_t) rows * 4),
-          bounds((size_t) (2 * (E + E / GROUP + 2)) * 4), grp_gu(GROUP * mmq::matrix_bytes(p.gu, 1280, N) + 4096),
+          bounds((size_t) (2 * (E + E / GROUP + 2)) * 4), grp_gu(GROUP * mmq::matrix_bytes(p.gu, FF2, N) + 4096),
           grp_d(GROUP * mmq::matrix_bytes(p.d, N, FF) + 4096) {}
 };
 // `direct`: no gathers - the products read the blobs where they are (they must lie at a fixed stride, every expert
@@ -188,7 +189,7 @@ struct MmqBufs {
 void run_mmq(const Pair& p, const Geo& geo, mmq::Context& ctx, MmqBufs& b, const Routing& r, const float* x_dev,
              const int32_t* src_dev, const std::vector<const uint8_t*>& blob, int64_t rows, cudaStream_t s,
              bool direct = false, size_t stride = 0) {
-    const size_t gub = mmq::matrix_bytes(p.gu, 1280, N), db = mmq::matrix_bytes(p.d, N, FF);
+    const size_t gub = mmq::matrix_bytes(p.gu, FF2, N), db = mmq::matrix_bytes(p.d, N, FF);
     const int E = (int) r.cnt.size();
     std::vector<int32_t> order;
     for (int e = 0; e < E; ++e) if (r.cnt[(size_t) e] > 0) order.push_back(e);
@@ -220,12 +221,12 @@ void run_mmq(const Pair& p, const Geo& geo, mmq::Context& ctx, MmqBufs& b, const
         }
         const uint8_t* b0 = blob[(size_t) order[j0]];
         mmq::Product gu;
-        gu.w = direct ? (const void*) b0 : b.grp_gu.p; gu.type = (int) p.gu; gu.w_rows = 1280; gu.w_cols = N;
+        gu.w = direct ? (const void*) b0 : b.grp_gu.p; gu.type = (int) p.gu; gu.w_rows = FF2; gu.w_cols = N;
         gu.expert_bytes = direct ? stride : gub;
         gu.n = ngx; gu.xq = b.xq.p; gu.bounds = b.bounds.as<int32_t>() + j0; gu.ids = b.ident.as<int32_t>();
-        gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = 1280;
+        gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = FF2;
         ctx.run(gu, s);
-        mmq::swiglu(b.gu.as<float>() + r0 * 1280, b.h.as<float>() + r0 * FF, nr, FF, false, s);
+        mmq::swiglu(b.gu.as<float>() + r0 * FF2, b.h.as<float>() + r0 * FF, nr, FF, false, s);
         mmq::quantize(b.h.as<float>() + r0 * FF, nullptr, b.hq.p, (int) p.d, FF, FF, nr, s);
         mmq::Product dn;
         dn.w = direct ? (const void*) (b0 + geo.down_off) : b.grp_d.p; dn.type = (int) p.d; dn.w_rows = N; dn.w_cols = FF;
@@ -430,7 +431,7 @@ void reference_part(const Pair& p, cudaStream_t s, int T = 256, int E = 64, int
 
 // part 2: one layer at a real chunk, timed
 void timing_part(const Pair& p, int T, cudaStream_t s) {
-    static const int E = [] { const char* v = std::getenv("S20_E"); return v != nullptr && std::atoi(v) > 0 ? std::atoi(v) : 512; }();   // S20_E=256: the Coder pack
+    static const int E = [] { const char* v = std::getenv("S20_E"); return v != nullptr && std::atoi(v) > 0 ? std::atoi(v) : (int) strata::core::kCompiledGeometry.n_expert; }();   // S20_E=256: the Coder pack
     const Geo geo(p);
     std::mt19937 rng(1360);
     const int64_t rows = (int64_t) T * K;
@@ -703,10 +704,12 @@ int main(int argc, char** argv) {
         for (int i = 1; i < argc; ++i)
             if (std::strncmp(argv[i], "--only=", 7) == 0) only = argv[i] + 7;
         bool timing = true, ref = true;   // --no-timing, --no-ref: skip a part
+        bool edge = false;                // --edge: the tails and tile boundaries (a handful of tokens, one token, 130 / 517 / 640) of the Q5_0 pairs, for every pair
         std::vector<int> chunks = {2048, 3584, 8192};   // --chunks=A,B: the timing part's chunk sizes
         for (int i = 1; i < argc; ++i) {
             if (std::strcmp(argv[i], "--no-timing") == 0) timing = false;
             if (std::strcmp(argv[i], "--no-ref") == 0) ref = false;
+            if (std::strcmp(argv[i], "--edge") == 0) edge = true;
             if (std::strncmp(argv[i], "--chunks=", 9) == 0) {
                 chunks.clear();
                 for (const char* c = argv[i] + 9; *c;) {
@@ -732,11 +735,12 @@ int main(int argc, char** argv) {
             }
             if (ref) reference_part(p, s);
             // Q5_0 down (round 0217): tails and boundaries - a handful of tokens, one token, a tile's edge (63, 65 rows an expert is not controlled, the counts follow the routing)
-            if (ref && p.d == GGML_TYPE_Q5_0) {
+            if (ref && (p.d == GGML_TYPE_Q5_0 || edge)) {
                 reference_part(p, s, 33, 16, 5);
                 reference_part(p, s, 1, 16, -1);
                 reference_part(p, s, 130, 32, 77);
                 reference_part(p, s, 517, 32, 300);
+                if (edge) reference_part(p, s, 640, 16, 300);   // five 64-row tiles an expert: the paired-tile (128-row) kernels
             }
         }
         if (timing) {

tests/cuda/prefill_fused_moe_test.cpp

@@ -32,9 +32,12 @@ namespace {
 namespace mmq = strata::prefill::mmq;
 namespace fused = strata::prefill::fused;
 
-constexpr int N = 2560, FF = 640, K = 10, GROUP = 16, Q2 = 42;   // GGML_TYPE_Q2_0
-constexpr size_t BLOB = 1382400, O_D_CODES = (size_t) 1280 * 640, O_GU_SC = O_D_CODES + (size_t) 2560 * 160,
-                 O_D_SC = O_GU_SC + (size_t) 1280 * 40 * 2;
+// The compiled geometry's widths and top-k (Flash-Next: 2560 x 640, top 10); the blob is the Strata Q2 layout at that shape
+// (Flash-Next: 1,382,400 B): gate/up codes [2 FF][N / 4 B], down codes [N][FF / 4 B], gate/up scales [2 FF][N / 64] f16, down scales [N][FF / 64] f16.
+constexpr int N = fused::kMoeEmbd, FF = fused::kMoeFF, FF2 = fused::kMoeGuRows, K = (int) strata::core::kCompiledGeometry.n_expert_used, GROUP = 16, Q2 = 42;   // GGML_TYPE_Q2_0
+constexpr int ROW_GU = N / 4, ROW_D = FF / 4, SC_GU = N / 64, SC_D = FF / 64;
+constexpr size_t O_D_CODES = (size_t) FF2 * ROW_GU, O_GU_SC = O_D_CODES + (size_t) N * ROW_D, O_D_SC = O_GU_SC + (size_t) FF2 * SC_GU * 2,
+                 BLOB = O_D_SC + (size_t) N * SC_D * 2;
 
 void ck(cudaError_t e, const char* what) {
     if (e != cudaSuccess) throw std::runtime_error(std::string(what) + ": " + cudaGetErrorString(e));
@@ -62,7 +65,7 @@ std::vector<uint8_t> make_blob(std::mt19937& rng) {
     return b;
 }
 
-// the blob's weights as floats: gate/up [1280][2560] (interleaved rows), down [2560][640]
+// the blob's weights as floats: gate/up [2 FF][N] (interleaved rows), down [N][FF]
 void dequant(const std::vector<uint8_t>& b, std::vector<float>& gu, std::vector<float>& dn) {
     auto row = [&](size_t codes, size_t ld, size_t scales, int nb, int r, float* out) {
         for (int k = 0; k < nb * 64; ++k) {
@@ -72,10 +75,10 @@ void dequant(const std::vector<uint8_t>& b, std::vector<float>& gu, std::vector<
             out[k] = ggml_fp16_to_fp32(h) * (float) (q - 1);
         }
     };
-    gu.resize((size_t) 1280 * N);
+    gu.resize((size_t) FF2 * N);
     dn.resize((size_t) N * FF);
-    for (int r = 0; r < 1280; ++r) row(0, 640, O_GU_SC, 40, r, gu.data() + (size_t) r * N);
-    for (int r = 0; r < N; ++r) row(O_D_CODES, 160, O_D_SC, 10, r, dn.data() + (size_t) r * FF);
+    for (int r = 0; r < FF2; ++r) row(0, ROW_GU, O_GU_SC, SC_GU, r, gu.data() + (size_t) r * N);
+    for (int r = 0; r < N; ++r) row(O_D_CODES, ROW_D, O_D_SC, SC_D, r, dn.data() + (size_t) r * FF);
 }
 
 // activations: N(0, 1) with a few large values (the hidden state has outliers); token `zero` all zero
@@ -132,14 +135,14 @@ Routing sort_rows(const std::vector<int32_t>& ids, int E) {
 struct MmqBufs {
     Dev xq, gu, h, hq, dm, ident, bounds, grp_gu, grp_d;
     MmqBufs(int64_t rows, int E)
-        : xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * 1280 * 4), h((size_t) rows * FF * 4),
+        : xq(mmq::q8_bytes(rows, N)), gu((size_t) rows * FF2 * 4), h((size_t) rows * FF * 4),
           hq(mmq::q8_bytes(rows, FF)),
           dm((size_t) rows * N * 4), ident((size_t) rows * 4), bounds((size_t) (2 * (E + E / GROUP + 2)) * 4),
-          grp_gu(GROUP * mmq::matrix_bytes(Q2, 1280, N) + 4096), grp_d(GROUP * mmq::matrix_bytes(Q2, N, FF) + 4096) {}
+          grp_gu(GROUP * mmq::matrix_bytes(Q2, FF2, N) + 4096), grp_d(GROUP * mmq::matrix_bytes(Q2, N, FF) + 4096) {}
 };
 void run_mmq(mmq::Context& ctx, MmqBufs& b, const Routing& r, const float* x_dev, const int32_t* src_dev,
              const std::vector<const uint8_t*>& blob, int64_t rows, cudaStream_t s) {
-    const size_t gub = mmq::matrix_bytes(Q2, 1280, N), db = mmq::matrix_bytes(Q2, N, FF);
+    const size_t gub = mmq::matrix_bytes(Q2, FF2, N), db = mmq::matrix_bytes(Q2, N, FF);
     const int E = (int) r.cnt.size();
     std::vector<int32_t> order;
     for (int e = 0; e < E; ++e) if (r.cnt[(size_t) e] > 0) order.push_back(e);
@@ -166,11 +169,11 @@ void run_mmq(mmq::Context& ctx, MmqBufs& b, const Routing& r, const float* x_dev
         ck(cudaMemsetAsync(b.grp_gu.as<uint8_t>() + (size_t) ngx * gub, 0, 4096, s), "tail");
         ck(cudaMemsetAsync(b.grp_d.as<uint8_t>() + (size_t) ngx * db, 0, 4096, s), "tail");
         mmq::Product gu;
-        gu.w = b.grp_gu.p; gu.type = Q2; gu.w_rows = 1280; gu.w_cols = N; gu.expert_bytes = gub;
+        gu.w = b.grp_gu.p; gu.type = Q2; gu.w_rows = FF2; gu.w_cols = N; gu.expert_bytes = gub;
         gu.n = ngx; gu.xq = b.xq.p; gu.bounds = b.bounds.as<int32_t>() + j0; gu.ids = b.ident.as<int32_t>();
-        gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = 1280;
+        gu.total_rows = rows; gu.max_rows = maxr; gu.dst = b.gu.as<float>(); gu.ld_dst = FF2;
         ctx.run(gu, s);
-        mmq::swiglu(b.gu.as<float>() + r0 * 1280, b.h.as<float>() + r0 * FF, nr, FF, true, s);
+        mmq::swiglu(b.gu.as<float>() + r0 * FF2, b.h.as<float>() + r0 * FF, nr, FF, true, s);
         mmq::quantize(b.h.as<float>() + r0 * FF, nullptr, b.hq.p, Q2, FF, FF, nr, s);
         mmq::Product dn;
         dn.w = b.grp_d.p; dn.type = Q2; dn.w_rows = N; dn.w_cols = FF; dn.expert_bytes = db;

tools/q4geo_synth.py

@@ -45,6 +45,13 @@ QTYPES = {
     "IQ4_XS": (23, 256, 136, [(0, 1e-5)]),
 }
 
+ROUTER_SCALE = 1.0   # --router-scale: the router's weights times this. Random weights give router logits packed so closely that two
+                     # implementations of one model (llama.cpp CPU and GPU among them) pick different experts at a few positions, and one
+                     # expert flip moves a position's KL by 100x; a larger scale spreads the logits like a trained router's
+EXACT_PARTS: set = set()   # --exact-parts router,f32,q8 (--exact-bf16 = all three): which tensors are written BF16-exact
+EXACT_BF16 = False   # --exact-bf16: every F32 tensor holds BF16-representable values and the Q8_0 scale is a power of two, so the pack's
+                     # --compat-bf16 conversion (router, shared-expert gate, HC down/up, ...) is lossless and llama.cpp sees the engine's weights
+
 # Flash-Next's values for every field a shape does not name
 FLASH = dict(n_embd=2560, n_layers=48, n_head=24, n_head_kv=2, head_dim=256, n_rot=64, ssm_state=128, ssm_k_heads=16,
              ssm_v_heads=48, ssm_d_conv=4, idx_q_heads=4, idx_key_dim=128, idx_top_k=2048, idx_block=4, hc=4, hc_lr=320,
@@ -157,10 +164,22 @@ def recipe_type(recipe: str, role: str, layer: int, s: dict) -> str:
         if role in ("gate", "up"):
             return "Q4_K"
         return "Q5_0"
+    if recipe == "fused-kq":         # Q4_K gate/up (Q5_K in layer 0) with Q5_1 / Q8_0 down at ANY expert width: the pairs the AMD fused / F16 prompt kernels take (Q4-HIPKQ)
+        if role in ("gate", "up"):
+            return "Q5_K" if first else "Q4_K"
+        return "Q8_0" if layer % 2 == 0 else "Q5_1"
+    if recipe == "mix-kq":           # one file with every K-quant expert path of the AMD prompt kernels: Q5_K/Q8_0 (fused), Q4_K/Q5_K and Q4_K/Q6_K (MMQ), Q4_K/Q5_1 (fused F16), by layer % 4 (Q4-HIPKQ)
+        if role in ("gate", "up"):
+            return "Q5_K" if layer % 4 == 0 else "Q4_K"
+        return ("Q8_0", "Q5_K", "Q5_1", "Q6_K")[layer % 4]
+    if recipe == "fused-iq":         # IQ3_S gate/up (IQ4_XS in layer 0) with IQ4_NL down at any expert width: the native IQ pack's fused prompt kernels (Q4-HIPKQ)
+        if role in ("gate", "up"):
+            return "IQ4_XS" if first else "IQ3_S"
+        return "IQ4_NL"
     raise SystemExit(f"unknown recipe {recipe}")
 
 
-RECIPES = ["ud-q4k", "ud-q4k-q6", "ud-iq4xs", "ud-q5_0down"]
+RECIPES = ["ud-q4k", "ud-q4k-q6", "ud-iq4xs", "ud-q5_0down", "fused-kq", "fused-iq", "mix-kq"]
 
 
 def row_ok(ty: str, ne0: int) -> bool:
@@ -168,7 +187,23 @@ def row_ok(ty: str, ne0: int) -> bool:
 
 
 # ------------------------------------------------------------------ data
+def bf16_round(a: np.ndarray) -> np.ndarray:
+    """float32 rounded to the nearest BF16 (ties to even), returned as float32."""
+    u = a.astype("<f4").view("<u4").astype(np.uint64)
+    u = (u + 0x7FFF + ((u >> 16) & 1)) & 0xFFFF0000
+    return u.astype("<u4").view("<f4")
+
+
 def small_f32(role: str, shape: list[int], rng: np.random.Generator) -> np.ndarray:
+    a = _small_f32(role, shape, rng)
+    if role == "router":
+        a = (a * np.float32(ROUTER_SCALE)).astype("<f4")
+    if role == "router":
+        return bf16_round(a) if ("router" in EXACT_PARTS) else a
+    return bf16_round(a) if ("f32" in EXACT_PARTS) else a
+
+
+def _small_f32(role: str, shape: list[int], rng: np.random.Generator) -> np.ndarray:
     n = int(np.prod(shape))
     if role == "norm":
         return (1.0 + 0.05 * rng.standard_normal(n)).astype("<f4")
@@ -211,6 +246,8 @@ def emit_data(fh, ty: str, role: str, shape: list[int], rng: np.random.Generator
         m = min(step, nblk - done)
         b = rng.integers(0, 256, size=(m, tsz), dtype=np.uint8)
         for off, val in fields:
+            if "q8" in EXACT_PARTS and ty == "Q8_0":
+                val = 2.0 ** -12
             b[:, off:off + 2] = np.frombuffer(np.float16(val).tobytes(), dtype=np.uint8)
         fh.write(b.tobytes())
         done += m
@@ -367,9 +404,21 @@ def main() -> int:
     ap.add_argument("--header", help="write the geometry header for the shape (or --from-gguf) here")
     ap.add_argument("--label", help="kGeometryName in the header (default: the shape, plus _<layers>l)")
     ap.add_argument("--from-gguf", help="read the shape from this file's metadata instead of --shape")
+    ap.add_argument("--exact-bf16", action="store_true", help="F32 tensors rounded to BF16 values and Q8_0 scale 2^-12: the pack's BF16 conversion is then lossless")
+    ap.add_argument("--router-scale", type=float, default=1.0, help="multiply the router weights (default 1: the files of Q4-GEO2)")
+    ap.add_argument("--exact-parts", default="", help="router, f32 (the other F32 tensors) and/or q8 (Q8_0 scale 2^-12): the BF16-exact parts of --exact-bf16")
     ap.add_argument("--list", action="store_true")
+    ap.add_argument("--small-ple", action="store_true",
+                    help="PLE head sizes of 16 small primes (a table of ~1.6M rows, ~50 MB) instead of Flash-Next's 320001536 rows (28 GB): the hash moduli are read from the file, so the engine runs it; "
+                         "for hosts where the sparse 28 GB table cannot be written (Windows)")
     ap.add_argument("--set", default="", help="comma list of field=value overrides of the shape (a bisection of one field), e.g. n_ff=1024,hc_lr=384")
     a = ap.parse_args()
+    global EXACT_BF16
+    EXACT_BF16 = a.exact_bf16
+    global EXACT_PARTS
+    EXACT_PARTS = {"router", "f32", "q8"} if a.exact_bf16 else set(x for x in a.exact_parts.split(",") if x)
+    global ROUTER_SCALE
+    ROUTER_SCALE = a.router_scale
     if a.list:
         for n in sorted(SHAPES):
             s = shape(n)
@@ -383,6 +432,15 @@ def main() -> int:
         s = shape(a.shape, a.layers, **{k: int(v) for k, v in (kv.split("=") for kv in a.set.split(",") if kv)})
         label = a.label or (a.shape + (f"_{a.layers}l" if a.layers else ""))
         src = f"shape {a.shape}" + (f", {a.layers} layers" if a.layers else "")
+    if a.small_ple:
+        global PLE_ROWS
+        primes, n = [], 100003
+        while len(primes) < 16:
+            if all(n % d for d in range(2, int(n ** 0.5) + 1)):
+                primes.append(n)
+            n += 2
+        PLE_VOCAB[:] = primes
+        PLE_ROWS = sum(primes)
     if a.header:
         pathlib.Path(a.header).write_text(header_text(s, label, src), encoding="utf8", newline="\n")
         print(f"wrote {a.header} ({label})")