Skip to content

Commit d0f7332

Browse files
userFRMclaude
authored and
Graf RAGov
committed
feat: cherry-pick PR danveloper#14 (MLX 8-bit dequant) onto nerds-odd-e base
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
1 parent 840aff8 commit d0f7332

2 files changed

Lines changed: 166 additions & 39 deletions

File tree

‎metal_infer/infer.m‎

Lines changed: 98 additions & 39 deletions
Original file line numberDiff line numberDiff line change
@@ -676,6 +676,44 @@ static void cpu_dequant_matvec(
676676
}
677677
}
678678

679+
// 8-bit dequantized matrix-vector multiply (CPU fallback)
680+
// W is stored as packed uint32 (4 x 8-bit values per uint32)
681+
// scales/biases are bfloat16 per group
682+
static void cpu_dequant_matvec_8bit(
683+
const uint32_t *W, const uint16_t *scales, const uint16_t *biases,
684+
const float *x, float *out,
685+
int out_dim, int in_dim, int group_size
686+
) {
687+
int num_groups = in_dim / group_size;
688+
int packed_per_group = group_size / 4;
689+
int packed_cols = in_dim / 4;
690+
691+
for (int row = 0; row < out_dim; row++) {
692+
float acc = 0.0f;
693+
const uint32_t *w_row = W + row * packed_cols;
694+
const uint16_t *s_row = scales + row * num_groups;
695+
const uint16_t *b_row = biases + row * num_groups;
696+
697+
for (int g = 0; g < num_groups; g++) {
698+
float scale = bf16_to_f32(s_row[g]);
699+
float bias = bf16_to_f32(b_row[g]);
700+
int base_packed = g * packed_per_group;
701+
int base_x = g * group_size;
702+
703+
for (int p = 0; p < packed_per_group; p++) {
704+
uint32_t packed = w_row[base_packed + p];
705+
int x_base = base_x + p * 4;
706+
707+
acc += ((float)((packed >> 0) & 0xFF) * scale + bias) * x[x_base + 0];
708+
acc += ((float)((packed >> 8) & 0xFF) * scale + bias) * x[x_base + 1];
709+
acc += ((float)((packed >> 16) & 0xFF) * scale + bias) * x[x_base + 2];
710+
acc += ((float)((packed >> 24) & 0xFF) * scale + bias) * x[x_base + 3];
711+
}
712+
}
713+
out[row] = acc;
714+
}
715+
}
716+
679717
// RMS normalization: out = x * w / rms(x)
680718
static void cpu_rms_norm(const float *x, const uint16_t *w_bf16, float *out, int dim, float eps) {
681719
float sum_sq = 0.0f;
@@ -847,6 +885,7 @@ static void cpu_conv1d_step(
847885
id<MTLComputePipelineState> matvec_v5; // LUT dequant variant
848886
id<MTLComputePipelineState> matvec_fast; // for in_dim > 4096
849887
id<MTLComputePipelineState> matvec_2bit; // 2-bit expert dequant kernel
888+
id<MTLComputePipelineState> matvec_8bit; // 8-bit gate dequant kernel
850889
id<MTLComputePipelineState> rms_norm_sum;
851890
id<MTLComputePipelineState> rms_norm_apply;
852891
id<MTLComputePipelineState> rms_norm_apply_bf16;
@@ -991,6 +1030,7 @@ static void cpu_conv1d_step(
9911030
ctx->matvec_v5 = makePipe(@"dequant_matvec_4bit_v5"); // LUT variant (no uint→float conversions)
9921031
ctx->matvec_fast = makePipe(@"dequant_matvec_4bit_fast");
9931032
ctx->matvec_2bit = makePipe(@"dequant_matvec_2bit");
1033+
ctx->matvec_8bit = makePipe(@"dequant_matvec_8bit");
9941034
ctx->rms_norm_sum = makePipe(@"rms_norm_sum_sq");
9951035
ctx->rms_norm_apply = makePipe(@"rms_norm_apply");
9961036
ctx->rms_norm_apply_bf16 = makePipe(@"rms_norm_apply_bf16");
@@ -1309,6 +1349,7 @@ static void fast_dequant_matvec(
13091349
uint32_t in_dim;
13101350
uint32_t group_size;
13111351
int batch_slot; // which batch_out[slot] to use for GPU output
1352+
int bits; // quantization bits: 4 (default) or 8 (gate weights)
13121353
} BatchMatvecSpec;
13131354

13141355
// Run N matmuls in a single command buffer. All share the same input vector.
@@ -1332,8 +1373,12 @@ static void gpu_batch_matvec(
13321373
id<MTLBuffer> o_buf = ctx->batch_out[s->batch_slot];
13331374

13341375
id<MTLComputeCommandEncoder> enc = [cmdbuf computeCommandEncoder];
1335-
int use_v3 = (s->in_dim <= 4096);
1336-
[enc setComputePipelineState: use_v3 ? ctx->matvec_v3 : ctx->matvec_fast];
1376+
if (s->bits == 8) {
1377+
[enc setComputePipelineState: ctx->matvec_8bit];
1378+
} else {
1379+
int use_v3 = (s->in_dim <= 4096);
1380+
[enc setComputePipelineState: use_v3 ? ctx->matvec_v3 : ctx->matvec_fast];
1381+
}
13371382
[enc setBuffer:ctx->wf_buf offset:w_off atIndex:0];
13381383
[enc setBuffer:ctx->wf_buf offset:s_off atIndex:1];
13391384
[enc setBuffer:ctx->wf_buf offset:b_off atIndex:2];
@@ -1343,7 +1388,7 @@ static void gpu_batch_matvec(
13431388
[enc setBytes:&s->in_dim length:4 atIndex:6];
13441389
[enc setBytes:&s->group_size length:4 atIndex:7];
13451390

1346-
if (use_v3) {
1391+
if (s->bits == 8 || s->in_dim <= 4096) {
13471392
uint32_t num_tgs = (s->out_dim + 7) / 8;
13481393
[enc dispatchThreadgroups:MTLSizeMake(num_tgs, 1, 1)
13491394
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
@@ -1386,8 +1431,12 @@ static void gpu_encode_batch_matvec(
13861431
id<MTLBuffer> o_buf = ctx->batch_out[s->batch_slot];
13871432

13881433
id<MTLComputeCommandEncoder> enc = [cmdbuf computeCommandEncoder];
1389-
int use_v3 = (s->in_dim <= 4096);
1390-
[enc setComputePipelineState: use_v3 ? ctx->matvec_v3 : ctx->matvec_fast];
1434+
if (s->bits == 8) {
1435+
[enc setComputePipelineState: ctx->matvec_8bit];
1436+
} else {
1437+
int use_v3 = (s->in_dim <= 4096);
1438+
[enc setComputePipelineState: use_v3 ? ctx->matvec_v3 : ctx->matvec_fast];
1439+
}
13911440
[enc setBuffer:ctx->wf_buf offset:w_off atIndex:0];
13921441
[enc setBuffer:ctx->wf_buf offset:s_off atIndex:1];
13931442
[enc setBuffer:ctx->wf_buf offset:b_off atIndex:2];
@@ -1397,7 +1446,7 @@ static void gpu_encode_batch_matvec(
13971446
[enc setBytes:&s->in_dim length:4 atIndex:6];
13981447
[enc setBytes:&s->group_size length:4 atIndex:7];
13991448

1400-
if (use_v3) {
1449+
if (s->bits == 8 || s->in_dim <= 4096) {
14011450
uint32_t num_tgs = (s->out_dim + 7) / 8;
14021451
[enc dispatchThreadgroups:MTLSizeMake(num_tgs, 1, 1)
14031452
threadsPerThreadgroup:MTLSizeMake(256, 1, 1)];
@@ -1846,8 +1895,13 @@ static void fast_batch_matvec(
18461895
} else {
18471896
for (int i = 0; i < num_specs; i++) {
18481897
BatchMatvecSpec *s = &specs[i];
1849-
cpu_dequant_matvec(s->W, s->scales, s->biases, x, s->out_cpu,
1850-
s->out_dim, s->in_dim, s->group_size);
1898+
if (s->bits == 8) {
1899+
cpu_dequant_matvec_8bit(s->W, s->scales, s->biases, x, s->out_cpu,
1900+
s->out_dim, s->in_dim, s->group_size);
1901+
} else {
1902+
cpu_dequant_matvec(s->W, s->scales, s->biases, x, s->out_cpu,
1903+
s->out_dim, s->in_dim, s->group_size);
1904+
}
18511905
}
18521906
}
18531907
}
@@ -2146,9 +2200,9 @@ static void full_attention_forward(
21462200
// Batch Q/K/V into one command buffer (3 dispatches, 1 commit)
21472201
if (qw && qs && qb && kw && ks && kb && vw && vs && vb) {
21482202
BatchMatvecSpec qkv_specs[3] = {
2149-
{ qw, qs, qb, q_proj_out, (uint32_t)q_proj_dim, HIDDEN_DIM, GROUP_SIZE, 0 },
2150-
{ kw, ks, kb, k, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 1 },
2151-
{ vw, vs, vb, v, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 2 },
2203+
{ qw, qs, qb, q_proj_out, (uint32_t)q_proj_dim, HIDDEN_DIM, GROUP_SIZE, 0, 4 },
2204+
{ kw, ks, kb, k, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 1, 4 },
2205+
{ vw, vs, vb, v, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 2, 4 },
21522206
};
21532207
fast_batch_matvec(normed, HIDDEN_DIM, qkv_specs, 3);
21542208
}
@@ -2415,10 +2469,10 @@ static void linear_attention_forward(
24152469
if (qkv_w && qkv_s && qkv_b && z_w && z_s && z_b &&
24162470
b_w && b_s && b_b && a_w && a_s && a_b) {
24172471
BatchMatvecSpec la_specs[4] = {
2418-
{ qkv_w, qkv_s, qkv_b, qkv, (uint32_t)qkv_dim, HIDDEN_DIM, GROUP_SIZE, 0 },
2419-
{ z_w, z_s, z_b, z, (uint32_t)z_dim, HIDDEN_DIM, GROUP_SIZE, 1 },
2420-
{ b_w, b_s, b_b, beta, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 2 },
2421-
{ a_w, a_s, a_b, alpha, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 3 },
2472+
{ qkv_w, qkv_s, qkv_b, qkv, (uint32_t)qkv_dim, HIDDEN_DIM, GROUP_SIZE, 0, 4 },
2473+
{ z_w, z_s, z_b, z, (uint32_t)z_dim, HIDDEN_DIM, GROUP_SIZE, 1, 4 },
2474+
{ b_w, b_s, b_b, beta, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 2, 4 },
2475+
{ a_w, a_s, a_b, alpha, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 3, 4 },
24222476
};
24232477
fast_batch_matvec(normed, HIDDEN_DIM, la_specs, 4);
24242478
}
@@ -2672,10 +2726,10 @@ static void moe_forward(
26722726
if (gate_w && gate_s && gate_b && sgw && sgs && sgb &&
26732727
suw && sus && sub && seg_w && seg_s && seg_b) {
26742728
BatchMatvecSpec moe_specs[4] = {
2675-
{ gate_w, gate_s, gate_b, gate_scores, (uint32_t)NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE, 0 },
2676-
{ sgw, sgs, sgb, shared_gate, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 1 },
2677-
{ suw, sus, sub, shared_up, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 2 },
2678-
{ seg_w, seg_s, seg_b, &shared_gate_score, 1, HIDDEN_DIM, GROUP_SIZE, 3 },
2729+
{ gate_w, gate_s, gate_b, gate_scores, (uint32_t)NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE, 0, 8 },
2730+
{ sgw, sgs, sgb, shared_gate, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 1, 4 },
2731+
{ suw, sus, sub, shared_up, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 2, 4 },
2732+
{ seg_w, seg_s, seg_b, &shared_gate_score, 1, HIDDEN_DIM, GROUP_SIZE, 3, 8 },
26792733
};
26802734
fast_batch_matvec(h_post, HIDDEN_DIM, moe_specs, 4);
26812735
}
@@ -4043,9 +4097,9 @@ static void fused_layer_forward(
40434097

40444098
if (lc->q_w && lc->q_s && lc->q_b && lc->k_w && lc->k_s && lc->k_b &&
40454099
lc->v_w && lc->v_s && lc->v_b) {
4046-
attn_specs[0] = (BatchMatvecSpec){ lc->q_w, lc->q_s, lc->q_b, q_proj_out, (uint32_t)q_proj_dim, HIDDEN_DIM, GROUP_SIZE, 0 };
4047-
attn_specs[1] = (BatchMatvecSpec){ lc->k_w, lc->k_s, lc->k_b, k_out, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 1 };
4048-
attn_specs[2] = (BatchMatvecSpec){ lc->v_w, lc->v_s, lc->v_b, v_out, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 2 };
4100+
attn_specs[0] = (BatchMatvecSpec){ lc->q_w, lc->q_s, lc->q_b, q_proj_out, (uint32_t)q_proj_dim, HIDDEN_DIM, GROUP_SIZE, 0, 4 };
4101+
attn_specs[1] = (BatchMatvecSpec){ lc->k_w, lc->k_s, lc->k_b, k_out, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 1, 4 };
4102+
attn_specs[2] = (BatchMatvecSpec){ lc->v_w, lc->v_s, lc->v_b, v_out, (uint32_t)kv_dim, HIDDEN_DIM, GROUP_SIZE, 2, 4 };
40494103
num_attn_specs = 3;
40504104
}
40514105
}
@@ -4061,10 +4115,10 @@ static void fused_layer_forward(
40614115

40624116
if (lc->qkv_w && lc->qkv_s && lc->qkv_b && lc->z_w && lc->z_s && lc->z_b &&
40634117
lc->b_w && lc->b_s && lc->b_b && lc->a_w && lc->a_s && lc->a_b) {
4064-
attn_specs[0] = (BatchMatvecSpec){ lc->qkv_w, lc->qkv_s, lc->qkv_b, qkv_out, (uint32_t)qkv_dim, HIDDEN_DIM, GROUP_SIZE, 0 };
4065-
attn_specs[1] = (BatchMatvecSpec){ lc->z_w, lc->z_s, lc->z_b, z_out, (uint32_t)z_dim, HIDDEN_DIM, GROUP_SIZE, 1 };
4066-
attn_specs[2] = (BatchMatvecSpec){ lc->b_w, lc->b_s, lc->b_b, beta_out, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 2 };
4067-
attn_specs[3] = (BatchMatvecSpec){ lc->a_w, lc->a_s, lc->a_b, alpha_out, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 3 };
4118+
attn_specs[0] = (BatchMatvecSpec){ lc->qkv_w, lc->qkv_s, lc->qkv_b, qkv_out, (uint32_t)qkv_dim, HIDDEN_DIM, GROUP_SIZE, 0, 4 };
4119+
attn_specs[1] = (BatchMatvecSpec){ lc->z_w, lc->z_s, lc->z_b, z_out, (uint32_t)z_dim, HIDDEN_DIM, GROUP_SIZE, 1, 4 };
4120+
attn_specs[2] = (BatchMatvecSpec){ lc->b_w, lc->b_s, lc->b_b, beta_out, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 2, 4 };
4121+
attn_specs[3] = (BatchMatvecSpec){ lc->a_w, lc->a_s, lc->a_b, alpha_out, (uint32_t)LINEAR_NUM_V_HEADS, HIDDEN_DIM, GROUP_SIZE, 3, 4 };
40684122
num_attn_specs = 4;
40694123
}
40704124
}
@@ -4359,8 +4413,13 @@ static void fused_layer_forward(
43594413
} else {
43604414
for (int i = 0; i < num_attn_specs; i++) {
43614415
BatchMatvecSpec *s = &attn_specs[i];
4362-
cpu_dequant_matvec(s->W, s->scales, s->biases, normed, s->out_cpu,
4363-
s->out_dim, s->in_dim, s->group_size);
4416+
if (s->bits == 8) {
4417+
cpu_dequant_matvec_8bit(s->W, s->scales, s->biases, normed, s->out_cpu,
4418+
s->out_dim, s->in_dim, s->group_size);
4419+
} else {
4420+
cpu_dequant_matvec(s->W, s->scales, s->biases, normed, s->out_cpu,
4421+
s->out_dim, s->in_dim, s->group_size);
4422+
}
43644423
}
43654424
}
43664425
if (g_timing_enabled) { t1 = now_ms(); g_timing.cmd1_submit += t1 - t0; }
@@ -4398,9 +4457,9 @@ static void fused_layer_forward(
43984457
memset(spec_scores, 0, NUM_EXPERTS * sizeof(float));
43994458

44004459
// Gate projection matvec on pre-attention normed input (CPU, ~0.1ms for 512x4096)
4401-
cpu_dequant_matvec(lc->gate_w, lc->gate_s, lc->gate_b,
4402-
normed, spec_scores,
4403-
NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE);
4460+
cpu_dequant_matvec_8bit(lc->gate_w, lc->gate_s, lc->gate_b,
4461+
normed, spec_scores,
4462+
NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE);
44044463
cpu_softmax(spec_scores, NUM_EXPERTS);
44054464

44064465
int spec_K = (K > MAX_K) ? MAX_K : K;
@@ -5013,10 +5072,10 @@ static void fused_layer_forward(
50135072
// ---- Enc 5+: routing + shared expert projections (read buf_input) ----
50145073
#if HAS_SHARED_EXPERT
50155074
BatchMatvecSpec moe_specs[4] = {
5016-
{ gate_w, gate_s, gate_b, gate_scores, (uint32_t)NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE, 0 },
5017-
{ sgw, sgs, sgb, shared_gate, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 1 },
5018-
{ suw, sus, sub, shared_up, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 2 },
5019-
{ seg_w, seg_s, seg_b, &shared_gate_score, 1, HIDDEN_DIM, GROUP_SIZE, 3 },
5075+
{ gate_w, gate_s, gate_b, gate_scores, (uint32_t)NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE, 0, 8 },
5076+
{ sgw, sgs, sgb, shared_gate, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 1, 4 },
5077+
{ suw, sus, sub, shared_up, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 2, 4 },
5078+
{ seg_w, seg_s, seg_b, &shared_gate_score, 1, HIDDEN_DIM, GROUP_SIZE, 3, 8 },
50205079
};
50215080
int num_moe_specs = 4;
50225081
#else
@@ -5070,10 +5129,10 @@ static void fused_layer_forward(
50705129
if (have_moe_weights) {
50715130
#if HAS_SHARED_EXPERT
50725131
BatchMatvecSpec moe_specs[4] = {
5073-
{ gate_w, gate_s, gate_b, gate_scores, (uint32_t)NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE, 0 },
5074-
{ sgw, sgs, sgb, shared_gate, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 1 },
5075-
{ suw, sus, sub, shared_up, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 2 },
5076-
{ seg_w, seg_s, seg_b, &shared_gate_score, 1, HIDDEN_DIM, GROUP_SIZE, 3 },
5132+
{ gate_w, gate_s, gate_b, gate_scores, (uint32_t)NUM_EXPERTS, HIDDEN_DIM, GROUP_SIZE, 0, 8 },
5133+
{ sgw, sgs, sgb, shared_gate, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 1, 4 },
5134+
{ suw, sus, sub, shared_up, (uint32_t)SHARED_INTERMEDIATE, HIDDEN_DIM, GROUP_SIZE, 2, 4 },
5135+
{ seg_w, seg_s, seg_b, &shared_gate_score, 1, HIDDEN_DIM, GROUP_SIZE, 3, 8 },
50775136
};
50785137
fast_batch_matvec(h_post, HIDDEN_DIM, moe_specs, 4);
50795138
#else

‎metal_infer/shaders.metal‎

Lines changed: 68 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -492,6 +492,74 @@ kernel void dequant_matvec_2bit(
492492
}
493493

494494

495+
// ============================================================================
496+
// Kernel 1e: 8-bit dequantized matrix-vector multiply (FMA-optimized)
497+
// ============================================================================
498+
// Same structure as dequant_matvec_4bit_v3 but for 8-bit quantization:
499+
// - 4 values per uint32 (vs 8 for 4-bit, 16 for 2-bit)
500+
// - packed_cols = in_dim / 4
501+
// - Extract bytes: (packed >> 0) & 0xFF, >> 8, >> 16, >> 24
502+
// - FMA-optimized: (byte * scale + bias) * x = fma(byte, scale*x, bias*x)
503+
// Used for gate routing weights and shared_expert_gate in Qwen3.5-397B.
504+
505+
kernel void dequant_matvec_8bit(
506+
device const uint32_t* W_packed [[buffer(0)]], // [out_dim, in_dim/4]
507+
device const uint16_t* scales [[buffer(1)]], // [out_dim, num_groups] bf16
508+
device const uint16_t* biases [[buffer(2)]], // [out_dim, num_groups] bf16
509+
device const float* x [[buffer(3)]], // [in_dim]
510+
device float* out [[buffer(4)]], // [out_dim]
511+
constant uint& out_dim [[buffer(5)]],
512+
constant uint& in_dim [[buffer(6)]],
513+
constant uint& group_size [[buffer(7)]],
514+
uint tgid [[threadgroup_position_in_grid]],
515+
uint lid [[thread_position_in_threadgroup]],
516+
uint simd_lane [[thread_index_in_simdgroup]],
517+
uint simd_group [[simdgroup_index_in_threadgroup]]
518+
) {
519+
uint row = tgid * ROWS_PER_TG + simd_group;
520+
uint packed_cols = in_dim / 4; // 4 values per uint32 for 8-bit
521+
uint num_groups = in_dim / group_size;
522+
523+
threadgroup float x_shared[4096];
524+
for (uint i = lid; i < in_dim; i += 256) {
525+
x_shared[i] = x[i];
526+
}
527+
threadgroup_barrier(mem_flags::mem_threadgroup);
528+
if (row >= out_dim) return;
529+
530+
device const uint32_t* w_row = W_packed + row * packed_cols;
531+
device const uint16_t* s_row = scales + row * num_groups;
532+
device const uint16_t* b_row = biases + row * num_groups;
533+
534+
float acc = 0.0f;
535+
536+
for (uint col = simd_lane; col < packed_cols; col += 32) {
537+
// group_size/4 packed words per group
538+
uint g = col / (group_size / 4);
539+
float scale = bf16_to_f32(s_row[g]);
540+
float bias = bf16_to_f32(b_row[g]);
541+
542+
uint32_t packed = w_row[col];
543+
uint x_base = col * 4;
544+
545+
float sx0 = scale * x_shared[x_base + 0]; float bx0 = bias * x_shared[x_base + 0];
546+
float sx1 = scale * x_shared[x_base + 1]; float bx1 = bias * x_shared[x_base + 1];
547+
float sx2 = scale * x_shared[x_base + 2]; float bx2 = bias * x_shared[x_base + 2];
548+
float sx3 = scale * x_shared[x_base + 3]; float bx3 = bias * x_shared[x_base + 3];
549+
550+
acc += fma(float((packed >> 0) & 0xFF), sx0, bx0);
551+
acc += fma(float((packed >> 8) & 0xFF), sx1, bx1);
552+
acc += fma(float((packed >> 16) & 0xFF), sx2, bx2);
553+
acc += fma(float((packed >> 24) & 0xFF), sx3, bx3);
554+
}
555+
556+
float sum = simd_sum(acc);
557+
if (simd_lane == 0) {
558+
out[row] = sum;
559+
}
560+
}
561+
562+
495563
// ============================================================================
496564
// Kernel 1d: FULLY OPTIMIZED with uint4 vector loads
497565
// ============================================================================

0 commit comments

Comments
 (0)