← Back to C-Kernel-Engine Docs Doxygen Source Documentation
 
Loading...
Searching...
No Matches
Functions
rmsnorm_qkv.c File Reference

Fused RMSNorm + QKV Projection. More...

#include <stdint.h>
#include <stddef.h>
#include <math.h>
#include <string.h>
#include "ckernel_quant.h"

Go to the source code of this file.

Functions

static float compute_rms_scale (const float *x, int n, float eps)
 
void rmsnorm_forward (const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
 
void rmsnorm_qkv_fp32_fused (const float *x, const float *rms_weight, const float *wq, const float *wk, const float *wv, float *q_out, float *k_out, float *v_out, int embed_dim, int q_dim, int kv_dim, float eps)
 
void rmsnorm_qkv_fp32_fused_v2 (const float *x, const float *rms_weight, const float *wq, const float *wk, const float *wv, float *q_out, float *k_out, float *v_out, int embed_dim, int q_dim, int kv_dim, float eps)
 
void rmsnorm_qkv_fp32_fused_v3 (const float *x, const float *rms_weight, const float *wq, const float *wk, const float *wv, float *q_out, float *k_out, float *v_out, int embed_dim, int q_dim, int kv_dim, float eps)
 
void rmsnorm_qkv_q4k_fused (const float *x, const float *rms_weight, const void *wq, const void *wk, const void *wv, float *q_out, float *k_out, float *v_out, int embed_dim, int q_dim, int kv_dim, float eps)
 
void rmsnorm_qkv_separate_fp32 (const float *x, const float *rms_weight, const float *wq, const float *wk, const float *wv, float *normed, float *q_out, float *k_out, float *v_out, int embed_dim, int q_dim, int kv_dim, float eps)
 

Detailed Description

Fused RMSNorm + QKV Projection.

CK-ENGINE KERNEL RULES:

  1. NO malloc/free - memory via bump allocator, pointers passed in
  2. NO OpenMP - parallelization at orchestrator/codegen layer
  3. API must define: inputs, outputs, workspace, and memory layouts
  4. Pure computation - deterministic, no side effects

After changes: make test && make llamacpp-parity-full

Part of C-Kernel-Engine v6.6 Fusion Kernels

PROBLEM: Non-fused version does 4 DRAM round-trips for 'normed' buffer: rmsnorm(x, weight, normed); // Write normed to DRAM gemv(wq, normed, q); // Read normed from DRAM gemv(wk, normed, k); // Read normed from DRAM gemv(wv, normed, v); // Read normed from DRAM

SOLUTION: Fused version keeps 'normed' in registers/L1, zero DRAM access: rmsnorm_qkv_fused(x, weight, wq, wk, wv, q, k, v);

EXPECTED SPEEDUP: 1.5-2x for this operation

Definition in file rmsnorm_qkv.c.

Function Documentation

◆ compute_rms_scale()

static float compute_rms_scale ( const float *  x,
int  n,
float  eps 
)
inlinestatic

Definition at line 54 of file rmsnorm_qkv.c.

54 {
55 float sum_sq = 0.0f;
56
57#ifdef __AVX2__
58 __m256 vsum = _mm256_setzero_ps();
59 int i = 0;
60 for (; i + 7 < n; i += 8) {
61 __m256 vx = _mm256_loadu_ps(x + i);
62 vsum = _mm256_fmadd_ps(vx, vx, vsum);
63 }
64 // Horizontal sum
65 __m128 vlow = _mm256_castps256_ps128(vsum);
66 __m128 vhigh = _mm256_extractf128_ps(vsum, 1);
67 vlow = _mm_add_ps(vlow, vhigh);
68 vlow = _mm_hadd_ps(vlow, vlow);
69 vlow = _mm_hadd_ps(vlow, vlow);
70 sum_sq = _mm_cvtss_f32(vlow);
71 // Remainder
72 for (; i < n; i++) {
73 sum_sq += x[i] * x[i];
74 }
75#else
76 for (int i = 0; i < n; i++) {
77 sum_sq += x[i] * x[i];
78 }
79#endif
80
81 float rms = sqrtf(sum_sq / (float)n + eps);
82 return 1.0f / rms;
83}

Referenced by rmsnorm_qkv_fp32_fused(), rmsnorm_qkv_fp32_fused_v2(), rmsnorm_qkv_fp32_fused_v3(), and rmsnorm_qkv_separate_fp32().

◆ rmsnorm_forward()

void rmsnorm_forward ( const float *  input,
const float *  gamma,
float *  output,
float *  rstd_cache,
int  tokens,
int  d_model,
int  aligned_embed_dim,
float  eps 
)
extern

Definition at line 621 of file rmsnorm_kernels.c.

629{
630 rmsnorm_forward_strided_f32(
631 input,
632 gamma,
633 output,
634 rstd_cache,
635 tokens,
636 d_model,
637 aligned_embed_dim,
638 aligned_embed_dim,
639 eps
640 );
641}
rmsnorm_forward_strided_f32
void rmsnorm_forward_strided_f32(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int input_stride, int output_stride, float eps)
Definition rmsnorm_kernels.c:492

Referenced by rmsnorm_qkv_q4k_fused().

◆ rmsnorm_qkv_fp32_fused()

void rmsnorm_qkv_fp32_fused ( const float *  x,
const float *  rms_weight,
const float *  wq,
const float *  wk,
const float *  wv,
float *  q_out,
float *  k_out,
float *  v_out,
int  embed_dim,
int  q_dim,
int  kv_dim,
float  eps 
)

Definition at line 89 of file rmsnorm_qkv.c.

102 {
103 /* Step 1: Compute RMS scale factor (stays in register) */
104 float scale = compute_rms_scale(x, embed_dim, eps);
105
106 /* Step 2: Fused normalize + project
107 *
108 * Key insight: We compute normed[i] = x[i] * rms_weight[i] * scale
109 * on-the-fly during the GEMV, never storing the full normed vector.
110 *
111 * For each output element:
112 * q[j] = sum_i( wq[j,i] * x[i] * rms_weight[i] * scale )
113 * = scale * sum_i( wq[j,i] * x[i] * rms_weight[i] )
114 */
115
116 /* Q projection */
117 for (int j = 0; j < q_dim; j++) {
118 float sum = 0.0f;
119 const float *wq_row = wq + j * embed_dim;
120
121#ifdef __AVX2__
122 __m256 vsum = _mm256_setzero_ps();
123 int i = 0;
124 for (; i + 7 < embed_dim; i += 8) {
125 __m256 vx = _mm256_loadu_ps(x + i);
126 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
127 __m256 vw = _mm256_loadu_ps(wq_row + i);
128 __m256 vnormed = _mm256_mul_ps(vx, vrms);
129 vsum = _mm256_fmadd_ps(vw, vnormed, vsum);
130 }
131 /* Horizontal sum */
132 __m128 vlow = _mm256_castps256_ps128(vsum);
133 __m128 vhigh = _mm256_extractf128_ps(vsum, 1);
134 vlow = _mm_add_ps(vlow, vhigh);
135 vlow = _mm_hadd_ps(vlow, vlow);
136 vlow = _mm_hadd_ps(vlow, vlow);
137 sum = _mm_cvtss_f32(vlow);
138 /* Remainder */
139 for (; i < embed_dim; i++) {
140 sum += wq_row[i] * x[i] * rms_weight[i];
141 }
142#else
143 for (int i = 0; i < embed_dim; i++) {
144 sum += wq_row[i] * x[i] * rms_weight[i];
145 }
146#endif
147 q_out[j] = sum * scale;
148 }
149
150 /* K projection */
151 for (int j = 0; j < kv_dim; j++) {
152 float sum = 0.0f;
153 const float *wk_row = wk + j * embed_dim;
154
155#ifdef __AVX2__
156 __m256 vsum = _mm256_setzero_ps();
157 int i = 0;
158 for (; i + 7 < embed_dim; i += 8) {
159 __m256 vx = _mm256_loadu_ps(x + i);
160 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
161 __m256 vw = _mm256_loadu_ps(wk_row + i);
162 __m256 vnormed = _mm256_mul_ps(vx, vrms);
163 vsum = _mm256_fmadd_ps(vw, vnormed, vsum);
164 }
165 __m128 vlow = _mm256_castps256_ps128(vsum);
166 __m128 vhigh = _mm256_extractf128_ps(vsum, 1);
167 vlow = _mm_add_ps(vlow, vhigh);
168 vlow = _mm_hadd_ps(vlow, vlow);
169 vlow = _mm_hadd_ps(vlow, vlow);
170 sum = _mm_cvtss_f32(vlow);
171 for (; i < embed_dim; i++) {
172 sum += wk_row[i] * x[i] * rms_weight[i];
173 }
174#else
175 for (int i = 0; i < embed_dim; i++) {
176 sum += wk_row[i] * x[i] * rms_weight[i];
177 }
178#endif
179 k_out[j] = sum * scale;
180 }
181
182 /* V projection */
183 for (int j = 0; j < kv_dim; j++) {
184 float sum = 0.0f;
185 const float *wv_row = wv + j * embed_dim;
186
187#ifdef __AVX2__
188 __m256 vsum = _mm256_setzero_ps();
189 int i = 0;
190 for (; i + 7 < embed_dim; i += 8) {
191 __m256 vx = _mm256_loadu_ps(x + i);
192 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
193 __m256 vw = _mm256_loadu_ps(wv_row + i);
194 __m256 vnormed = _mm256_mul_ps(vx, vrms);
195 vsum = _mm256_fmadd_ps(vw, vnormed, vsum);
196 }
197 __m128 vlow = _mm256_castps256_ps128(vsum);
198 __m128 vhigh = _mm256_extractf128_ps(vsum, 1);
199 vlow = _mm_add_ps(vlow, vhigh);
200 vlow = _mm_hadd_ps(vlow, vlow);
201 vlow = _mm_hadd_ps(vlow, vlow);
202 sum = _mm_cvtss_f32(vlow);
203 for (; i < embed_dim; i++) {
204 sum += wv_row[i] * x[i] * rms_weight[i];
205 }
206#else
207 for (int i = 0; i < embed_dim; i++) {
208 sum += wv_row[i] * x[i] * rms_weight[i];
209 }
210#endif
211 v_out[j] = sum * scale;
212 }
213}
compute_rms_scale
static float compute_rms_scale(const float *x, int n, float eps)
Definition rmsnorm_qkv.c:54

References compute_rms_scale().

◆ rmsnorm_qkv_fp32_fused_v2()

void rmsnorm_qkv_fp32_fused_v2 ( const float *  x,
const float *  rms_weight,
const float *  wq,
const float *  wk,
const float *  wv,
float *  q_out,
float *  k_out,
float *  v_out,
int  embed_dim,
int  q_dim,
int  kv_dim,
float  eps 
)

Definition at line 316 of file rmsnorm_qkv.c.

329 {
330 /* Step 1: Compute RMS scale (requires full pass - unavoidable) */
331 float scale = compute_rms_scale(x, embed_dim, eps);
332
333#ifdef __AVX2__
334 __m256 vscale = _mm256_set1_ps(scale);
335
336 /* ═══════════════════════════════════════════════════════════════════════
337 * Q PROJECTION: Process 8 outputs at a time (one cache line)
338 *
339 * For each output cache line [j:j+8]:
340 * acc[0..7] = 0
341 * For each input cache line [i:i+8]:
342 * normed = x[i:i+8] * rms_weight[i:i+8] * scale ← IN REGISTER!
343 * acc[k] += W[j+k, i:i+8] · normed ← FMADD
344 * Store q_out[j:j+8]
345 * ═══════════════════════════════════════════════════════════════════════ */
346
347 for (int j = 0; j < q_dim; j += 8) {
348 /* 8 accumulators for 8 output elements - all in YMM registers */
349 __m256 acc0 = _mm256_setzero_ps();
350 __m256 acc1 = _mm256_setzero_ps();
351 __m256 acc2 = _mm256_setzero_ps();
352 __m256 acc3 = _mm256_setzero_ps();
353 __m256 acc4 = _mm256_setzero_ps();
354 __m256 acc5 = _mm256_setzero_ps();
355 __m256 acc6 = _mm256_setzero_ps();
356 __m256 acc7 = _mm256_setzero_ps();
357
358 /* Process input in cache-line chunks */
359 int i = 0;
360 for (; i + 7 < embed_dim; i += 8) {
361 /* Load input cache line and normalize IN REGISTER */
362 __m256 vx = _mm256_loadu_ps(x + i);
363 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
364 __m256 normed = _mm256_mul_ps(_mm256_mul_ps(vx, vrms), vscale);
365 /* normed is now in YMM register - NEVER touches memory! */
366
367 /* Load 8 weight rows and accumulate */
368 /* Each row is at wq[(j+k)*embed_dim + i] */
369 if (j + 0 < q_dim) {
370 __m256 w0 = _mm256_loadu_ps(wq + (j+0)*embed_dim + i);
371 acc0 = _mm256_fmadd_ps(w0, normed, acc0);
372 }
373 if (j + 1 < q_dim) {
374 __m256 w1 = _mm256_loadu_ps(wq + (j+1)*embed_dim + i);
375 acc1 = _mm256_fmadd_ps(w1, normed, acc1);
376 }
377 if (j + 2 < q_dim) {
378 __m256 w2 = _mm256_loadu_ps(wq + (j+2)*embed_dim + i);
379 acc2 = _mm256_fmadd_ps(w2, normed, acc2);
380 }
381 if (j + 3 < q_dim) {
382 __m256 w3 = _mm256_loadu_ps(wq + (j+3)*embed_dim + i);
383 acc3 = _mm256_fmadd_ps(w3, normed, acc3);
384 }
385 if (j + 4 < q_dim) {
386 __m256 w4 = _mm256_loadu_ps(wq + (j+4)*embed_dim + i);
387 acc4 = _mm256_fmadd_ps(w4, normed, acc4);
388 }
389 if (j + 5 < q_dim) {
390 __m256 w5 = _mm256_loadu_ps(wq + (j+5)*embed_dim + i);
391 acc5 = _mm256_fmadd_ps(w5, normed, acc5);
392 }
393 if (j + 6 < q_dim) {
394 __m256 w6 = _mm256_loadu_ps(wq + (j+6)*embed_dim + i);
395 acc6 = _mm256_fmadd_ps(w6, normed, acc6);
396 }
397 if (j + 7 < q_dim) {
398 __m256 w7 = _mm256_loadu_ps(wq + (j+7)*embed_dim + i);
399 acc7 = _mm256_fmadd_ps(w7, normed, acc7);
400 }
401 }
402
403 /* Handle remainder (scalar, rare for aligned dims) */
404 for (; i < embed_dim; i++) {
405 float normed_scalar = x[i] * rms_weight[i] * scale;
406 if (j + 0 < q_dim) acc0 = _mm256_add_ps(acc0, _mm256_set1_ps(wq[(j+0)*embed_dim + i] * normed_scalar));
407 if (j + 1 < q_dim) acc1 = _mm256_add_ps(acc1, _mm256_set1_ps(wq[(j+1)*embed_dim + i] * normed_scalar));
408 if (j + 2 < q_dim) acc2 = _mm256_add_ps(acc2, _mm256_set1_ps(wq[(j+2)*embed_dim + i] * normed_scalar));
409 if (j + 3 < q_dim) acc3 = _mm256_add_ps(acc3, _mm256_set1_ps(wq[(j+3)*embed_dim + i] * normed_scalar));
410 if (j + 4 < q_dim) acc4 = _mm256_add_ps(acc4, _mm256_set1_ps(wq[(j+4)*embed_dim + i] * normed_scalar));
411 if (j + 5 < q_dim) acc5 = _mm256_add_ps(acc5, _mm256_set1_ps(wq[(j+5)*embed_dim + i] * normed_scalar));
412 if (j + 6 < q_dim) acc6 = _mm256_add_ps(acc6, _mm256_set1_ps(wq[(j+6)*embed_dim + i] * normed_scalar));
413 if (j + 7 < q_dim) acc7 = _mm256_add_ps(acc7, _mm256_set1_ps(wq[(j+7)*embed_dim + i] * normed_scalar));
414 }
415
416 /* Horizontal sum and store output cache line */
417 if (j + 0 < q_dim) q_out[j+0] = hsum256_ps(acc0);
418 if (j + 1 < q_dim) q_out[j+1] = hsum256_ps(acc1);
419 if (j + 2 < q_dim) q_out[j+2] = hsum256_ps(acc2);
420 if (j + 3 < q_dim) q_out[j+3] = hsum256_ps(acc3);
421 if (j + 4 < q_dim) q_out[j+4] = hsum256_ps(acc4);
422 if (j + 5 < q_dim) q_out[j+5] = hsum256_ps(acc5);
423 if (j + 6 < q_dim) q_out[j+6] = hsum256_ps(acc6);
424 if (j + 7 < q_dim) q_out[j+7] = hsum256_ps(acc7);
425 }
426
427 /* ═══════════════════════════════════════════════════════════════════════
428 * K PROJECTION: Same pattern, smaller output
429 * ═══════════════════════════════════════════════════════════════════════ */
430
431 for (int j = 0; j < kv_dim; j += 8) {
432 __m256 acc0 = _mm256_setzero_ps();
433 __m256 acc1 = _mm256_setzero_ps();
434 __m256 acc2 = _mm256_setzero_ps();
435 __m256 acc3 = _mm256_setzero_ps();
436 __m256 acc4 = _mm256_setzero_ps();
437 __m256 acc5 = _mm256_setzero_ps();
438 __m256 acc6 = _mm256_setzero_ps();
439 __m256 acc7 = _mm256_setzero_ps();
440
441 int i = 0;
442 for (; i + 7 < embed_dim; i += 8) {
443 __m256 vx = _mm256_loadu_ps(x + i);
444 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
445 __m256 normed = _mm256_mul_ps(_mm256_mul_ps(vx, vrms), vscale);
446
447 if (j + 0 < kv_dim) acc0 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+0)*embed_dim + i), normed, acc0);
448 if (j + 1 < kv_dim) acc1 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+1)*embed_dim + i), normed, acc1);
449 if (j + 2 < kv_dim) acc2 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+2)*embed_dim + i), normed, acc2);
450 if (j + 3 < kv_dim) acc3 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+3)*embed_dim + i), normed, acc3);
451 if (j + 4 < kv_dim) acc4 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+4)*embed_dim + i), normed, acc4);
452 if (j + 5 < kv_dim) acc5 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+5)*embed_dim + i), normed, acc5);
453 if (j + 6 < kv_dim) acc6 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+6)*embed_dim + i), normed, acc6);
454 if (j + 7 < kv_dim) acc7 = _mm256_fmadd_ps(_mm256_loadu_ps(wk + (j+7)*embed_dim + i), normed, acc7);
455 }
456
457 for (; i < embed_dim; i++) {
458 float normed_scalar = x[i] * rms_weight[i] * scale;
459 if (j + 0 < kv_dim) acc0 = _mm256_add_ps(acc0, _mm256_set1_ps(wk[(j+0)*embed_dim + i] * normed_scalar));
460 if (j + 1 < kv_dim) acc1 = _mm256_add_ps(acc1, _mm256_set1_ps(wk[(j+1)*embed_dim + i] * normed_scalar));
461 if (j + 2 < kv_dim) acc2 = _mm256_add_ps(acc2, _mm256_set1_ps(wk[(j+2)*embed_dim + i] * normed_scalar));
462 if (j + 3 < kv_dim) acc3 = _mm256_add_ps(acc3, _mm256_set1_ps(wk[(j+3)*embed_dim + i] * normed_scalar));
463 if (j + 4 < kv_dim) acc4 = _mm256_add_ps(acc4, _mm256_set1_ps(wk[(j+4)*embed_dim + i] * normed_scalar));
464 if (j + 5 < kv_dim) acc5 = _mm256_add_ps(acc5, _mm256_set1_ps(wk[(j+5)*embed_dim + i] * normed_scalar));
465 if (j + 6 < kv_dim) acc6 = _mm256_add_ps(acc6, _mm256_set1_ps(wk[(j+6)*embed_dim + i] * normed_scalar));
466 if (j + 7 < kv_dim) acc7 = _mm256_add_ps(acc7, _mm256_set1_ps(wk[(j+7)*embed_dim + i] * normed_scalar));
467 }
468
469 if (j + 0 < kv_dim) k_out[j+0] = hsum256_ps(acc0);
470 if (j + 1 < kv_dim) k_out[j+1] = hsum256_ps(acc1);
471 if (j + 2 < kv_dim) k_out[j+2] = hsum256_ps(acc2);
472 if (j + 3 < kv_dim) k_out[j+3] = hsum256_ps(acc3);
473 if (j + 4 < kv_dim) k_out[j+4] = hsum256_ps(acc4);
474 if (j + 5 < kv_dim) k_out[j+5] = hsum256_ps(acc5);
475 if (j + 6 < kv_dim) k_out[j+6] = hsum256_ps(acc6);
476 if (j + 7 < kv_dim) k_out[j+7] = hsum256_ps(acc7);
477 }
478
479 /* ═══════════════════════════════════════════════════════════════════════
480 * V PROJECTION: Same pattern
481 * ═══════════════════════════════════════════════════════════════════════ */
482
483 for (int j = 0; j < kv_dim; j += 8) {
484 __m256 acc0 = _mm256_setzero_ps();
485 __m256 acc1 = _mm256_setzero_ps();
486 __m256 acc2 = _mm256_setzero_ps();
487 __m256 acc3 = _mm256_setzero_ps();
488 __m256 acc4 = _mm256_setzero_ps();
489 __m256 acc5 = _mm256_setzero_ps();
490 __m256 acc6 = _mm256_setzero_ps();
491 __m256 acc7 = _mm256_setzero_ps();
492
493 int i = 0;
494 for (; i + 7 < embed_dim; i += 8) {
495 __m256 vx = _mm256_loadu_ps(x + i);
496 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
497 __m256 normed = _mm256_mul_ps(_mm256_mul_ps(vx, vrms), vscale);
498
499 if (j + 0 < kv_dim) acc0 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+0)*embed_dim + i), normed, acc0);
500 if (j + 1 < kv_dim) acc1 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+1)*embed_dim + i), normed, acc1);
501 if (j + 2 < kv_dim) acc2 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+2)*embed_dim + i), normed, acc2);
502 if (j + 3 < kv_dim) acc3 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+3)*embed_dim + i), normed, acc3);
503 if (j + 4 < kv_dim) acc4 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+4)*embed_dim + i), normed, acc4);
504 if (j + 5 < kv_dim) acc5 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+5)*embed_dim + i), normed, acc5);
505 if (j + 6 < kv_dim) acc6 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+6)*embed_dim + i), normed, acc6);
506 if (j + 7 < kv_dim) acc7 = _mm256_fmadd_ps(_mm256_loadu_ps(wv + (j+7)*embed_dim + i), normed, acc7);
507 }
508
509 for (; i < embed_dim; i++) {
510 float normed_scalar = x[i] * rms_weight[i] * scale;
511 if (j + 0 < kv_dim) acc0 = _mm256_add_ps(acc0, _mm256_set1_ps(wv[(j+0)*embed_dim + i] * normed_scalar));
512 if (j + 1 < kv_dim) acc1 = _mm256_add_ps(acc1, _mm256_set1_ps(wv[(j+1)*embed_dim + i] * normed_scalar));
513 if (j + 2 < kv_dim) acc2 = _mm256_add_ps(acc2, _mm256_set1_ps(wv[(j+2)*embed_dim + i] * normed_scalar));
514 if (j + 3 < kv_dim) acc3 = _mm256_add_ps(acc3, _mm256_set1_ps(wv[(j+3)*embed_dim + i] * normed_scalar));
515 if (j + 4 < kv_dim) acc4 = _mm256_add_ps(acc4, _mm256_set1_ps(wv[(j+4)*embed_dim + i] * normed_scalar));
516 if (j + 5 < kv_dim) acc5 = _mm256_add_ps(acc5, _mm256_set1_ps(wv[(j+5)*embed_dim + i] * normed_scalar));
517 if (j + 6 < kv_dim) acc6 = _mm256_add_ps(acc6, _mm256_set1_ps(wv[(j+6)*embed_dim + i] * normed_scalar));
518 if (j + 7 < kv_dim) acc7 = _mm256_add_ps(acc7, _mm256_set1_ps(wv[(j+7)*embed_dim + i] * normed_scalar));
519 }
520
521 if (j + 0 < kv_dim) v_out[j+0] = hsum256_ps(acc0);
522 if (j + 1 < kv_dim) v_out[j+1] = hsum256_ps(acc1);
523 if (j + 2 < kv_dim) v_out[j+2] = hsum256_ps(acc2);
524 if (j + 3 < kv_dim) v_out[j+3] = hsum256_ps(acc3);
525 if (j + 4 < kv_dim) v_out[j+4] = hsum256_ps(acc4);
526 if (j + 5 < kv_dim) v_out[j+5] = hsum256_ps(acc5);
527 if (j + 6 < kv_dim) v_out[j+6] = hsum256_ps(acc6);
528 if (j + 7 < kv_dim) v_out[j+7] = hsum256_ps(acc7);
529 }
530
531#else
532 /* Scalar fallback - same logic, no SIMD */
533 for (int j = 0; j < q_dim; j++) {
534 float sum = 0.0f;
535 for (int i = 0; i < embed_dim; i++) {
536 float normed = x[i] * rms_weight[i] * scale;
537 sum += wq[j * embed_dim + i] * normed;
538 }
539 q_out[j] = sum;
540 }
541 for (int j = 0; j < kv_dim; j++) {
542 float sum = 0.0f;
543 for (int i = 0; i < embed_dim; i++) {
544 float normed = x[i] * rms_weight[i] * scale;
545 sum += wk[j * embed_dim + i] * normed;
546 }
547 k_out[j] = sum;
548 }
549 for (int j = 0; j < kv_dim; j++) {
550 float sum = 0.0f;
551 for (int i = 0; i < embed_dim; i++) {
552 float normed = x[i] * rms_weight[i] * scale;
553 sum += wv[j * embed_dim + i] * normed;
554 }
555 v_out[j] = sum;
556 }
557#endif
558}

References compute_rms_scale().

◆ rmsnorm_qkv_fp32_fused_v3()

void rmsnorm_qkv_fp32_fused_v3 ( const float *  x,
const float *  rms_weight,
const float *  wq,
const float *  wk,
const float *  wv,
float *  q_out,
float *  k_out,
float *  v_out,
int  embed_dim,
int  q_dim,
int  kv_dim,
float  eps 
)

Definition at line 586 of file rmsnorm_qkv.c.

599 {
600 /* Step 1: Compute RMS scale (requires full pass - unavoidable) */
601 float scale = compute_rms_scale(x, embed_dim, eps);
602
603#ifdef __AVX2__
604 __m256 vscale = _mm256_set1_ps(scale);
605
606 /* ═══════════════════════════════════════════════════════════════════════
607 * Phase 1: Process Q outputs that have corresponding K,V outputs
608 * (j < kv_dim: compute Q, K, V together)
609 * ═══════════════════════════════════════════════════════════════════════ */
610 for (int j = 0; j < kv_dim; j++) {
611 __m256 q_acc = _mm256_setzero_ps();
612 __m256 k_acc = _mm256_setzero_ps();
613 __m256 v_acc = _mm256_setzero_ps();
614
615 int i = 0;
616 for (; i + 7 < embed_dim; i += 8) {
617 /* Load input and normalize - computed ONCE, used THREE times! */
618 __m256 vx = _mm256_loadu_ps(x + i);
619 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
620 __m256 normed = _mm256_mul_ps(_mm256_mul_ps(vx, vrms), vscale);
621
622 /* Load weight rows - sequential access per row */
623 __m256 wq_row = _mm256_loadu_ps(wq + j * embed_dim + i);
624 __m256 wk_row = _mm256_loadu_ps(wk + j * embed_dim + i);
625 __m256 wv_row = _mm256_loadu_ps(wv + j * embed_dim + i);
626
627 /* Accumulate - normed stays in register! */
628 q_acc = _mm256_fmadd_ps(wq_row, normed, q_acc);
629 k_acc = _mm256_fmadd_ps(wk_row, normed, k_acc);
630 v_acc = _mm256_fmadd_ps(wv_row, normed, v_acc);
631 }
632
633 /* Handle remainder (scalar) */
634 float q_sum = hsum256_ps(q_acc);
635 float k_sum = hsum256_ps(k_acc);
636 float v_sum = hsum256_ps(v_acc);
637
638 for (; i < embed_dim; i++) {
639 float normed = x[i] * rms_weight[i] * scale;
640 q_sum += wq[j * embed_dim + i] * normed;
641 k_sum += wk[j * embed_dim + i] * normed;
642 v_sum += wv[j * embed_dim + i] * normed;
643 }
644
645 q_out[j] = q_sum;
646 k_out[j] = k_sum;
647 v_out[j] = v_sum;
648 }
649
650 /* ═══════════════════════════════════════════════════════════════════════
651 * Phase 2: Process remaining Q outputs (j >= kv_dim: Q only)
652 * This handles GQA where q_dim > kv_dim
653 * ═══════════════════════════════════════════════════════════════════════ */
654 for (int j = kv_dim; j < q_dim; j++) {
655 __m256 q_acc = _mm256_setzero_ps();
656
657 int i = 0;
658 for (; i + 7 < embed_dim; i += 8) {
659 __m256 vx = _mm256_loadu_ps(x + i);
660 __m256 vrms = _mm256_loadu_ps(rms_weight + i);
661 __m256 normed = _mm256_mul_ps(_mm256_mul_ps(vx, vrms), vscale);
662
663 __m256 wq_row = _mm256_loadu_ps(wq + j * embed_dim + i);
664 q_acc = _mm256_fmadd_ps(wq_row, normed, q_acc);
665 }
666
667 float q_sum = hsum256_ps(q_acc);
668 for (; i < embed_dim; i++) {
669 float normed = x[i] * rms_weight[i] * scale;
670 q_sum += wq[j * embed_dim + i] * normed;
671 }
672
673 q_out[j] = q_sum;
674 }
675
676#else
677 /* Scalar fallback - same simultaneous Q,K,V approach */
678 for (int j = 0; j < kv_dim; j++) {
679 float q_sum = 0.0f, k_sum = 0.0f, v_sum = 0.0f;
680 for (int i = 0; i < embed_dim; i++) {
681 float normed = x[i] * rms_weight[i] * scale;
682 q_sum += wq[j * embed_dim + i] * normed;
683 k_sum += wk[j * embed_dim + i] * normed;
684 v_sum += wv[j * embed_dim + i] * normed;
685 }
686 q_out[j] = q_sum;
687 k_out[j] = k_sum;
688 v_out[j] = v_sum;
689 }
690 for (int j = kv_dim; j < q_dim; j++) {
691 float q_sum = 0.0f;
692 for (int i = 0; i < embed_dim; i++) {
693 float normed = x[i] * rms_weight[i] * scale;
694 q_sum += wq[j * embed_dim + i] * normed;
695 }
696 q_out[j] = q_sum;
697 }
698#endif
699}

References compute_rms_scale().

◆ rmsnorm_qkv_q4k_fused()

void rmsnorm_qkv_q4k_fused ( const float *  x,
const float *  rms_weight,
const void *  wq,
const void *  wk,
const void *  wv,
float *  q_out,
float *  k_out,
float *  v_out,
int  embed_dim,
int  q_dim,
int  kv_dim,
float  eps 
)

Definition at line 222 of file rmsnorm_qkv.c.

235 {
236 /* Compute normalized input once using the same tuned kernel as the
237 * separate baseline. This keeps the fused-vs-separate comparison focused
238 * on avoiding extra call boundaries and keeping normed[] hot for the 3 GEMVs.
239 *
240 * For Q4_K, we can't easily fuse the normalization into the dequant loop
241 * because the block structure is complex. So we compute normed[] first,
242 * but keep it small enough to fit in L1 cache.
243 *
244 * TODO: For maximum performance, implement a true fused Q4_K GEMV
245 * that dequantizes and multiplies by normed[i] in the same loop.
246 */
247
248 /* Allocate normed on stack (fits in L1 for typical embed_dim <= 4096). */
249#if defined(__GNUC__) || defined(__clang__)
250 __attribute__((aligned(64))) float normed[4096];
251#else
252 float normed[4096];
253#endif
254 if (embed_dim > 4096) {
255 /* Fallback for very large models */
256 return; /* TODO: heap allocation */
257 }
258
259 rmsnorm_forward(x, rms_weight, normed, NULL, 1, embed_dim, embed_dim, eps);
260
261 /* Step 3: Q4_K GEMV with normed input
262 *
263 * Call existing Q4_K GEMV kernel with normed[] as input.
264 * The normed[] buffer is in L1 cache, so this is still fast.
265 *
266 * Key insight: normed[] never leaves L1 cache because we use it
267 * immediately in the next 3 GEMVs. This eliminates the DRAM write
268 * that would happen in the non-fused version.
269 */
270
271 /* Declare external GEMV function */
272 extern void gemv_q4_k(float *y, const void *W, const float *x, int M, int K);
273
274 /* Q projection: q_out[q_dim] = wq[q_dim, embed_dim] @ normed[embed_dim] */
275 gemv_q4_k(q_out, wq, normed, q_dim, embed_dim);
276
277 /* K projection: k_out[kv_dim] = wk[kv_dim, embed_dim] @ normed[embed_dim] */
278 gemv_q4_k(k_out, wk, normed, kv_dim, embed_dim);
279
280 /* V projection: v_out[kv_dim] = wv[kv_dim, embed_dim] @ normed[embed_dim] */
281 gemv_q4_k(v_out, wv, normed, kv_dim, embed_dim);
282}
gemv_q4_k
void gemv_q4_k(float *y, const void *W, const float *x, int M, int K)
Auto-dispatch GEMV based on available SIMD.
Definition gemm_kernels_q4k.c:301
rmsnorm_forward
void rmsnorm_forward(const float *input, const float *gamma, float *output, float *rstd_cache, int tokens, int d_model, int aligned_embed_dim, float eps)
Definition rmsnorm_kernels.c:621
__attribute__
__attribute__((visibility("default"))) CKTokenizer *ck_tokenizer_create(CKTokenizerType type)

References __attribute__(), gemv_q4_k(), and rmsnorm_forward().

◆ rmsnorm_qkv_separate_fp32()

void rmsnorm_qkv_separate_fp32 ( const float *  x,
const float *  rms_weight,
const float *  wq,
const float *  wk,
const float *  wv,
float *  normed,
float *  q_out,
float *  k_out,
float *  v_out,
int  embed_dim,
int  q_dim,
int  kv_dim,
float  eps 
)

Definition at line 707 of file rmsnorm_qkv.c.

721 {
722 /* Step 1: RMSNorm - writes normed to DRAM */
723 float scale = compute_rms_scale(x, embed_dim, eps);
724 for (int i = 0; i < embed_dim; i++) {
725 normed[i] = x[i] * rms_weight[i] * scale;
726 }
727
728 /* Step 2: Q projection - reads normed from DRAM */
729 for (int j = 0; j < q_dim; j++) {
730 float sum = 0.0f;
731 for (int i = 0; i < embed_dim; i++) {
732 sum += wq[j * embed_dim + i] * normed[i];
733 }
734 q_out[j] = sum;
735 }
736
737 /* Step 3: K projection - reads normed from DRAM */
738 for (int j = 0; j < kv_dim; j++) {
739 float sum = 0.0f;
740 for (int i = 0; i < embed_dim; i++) {
741 sum += wk[j * embed_dim + i] * normed[i];
742 }
743 k_out[j] = sum;
744 }
745
746 /* Step 4: V projection - reads normed from DRAM */
747 for (int j = 0; j < kv_dim; j++) {
748 float sum = 0.0f;
749 for (int i = 0; i < embed_dim; i++) {
750 sum += wv[j * embed_dim + i] * normed[i];
751 }
752 v_out[j] = sum;
753 }
754}

References compute_rms_scale().