Kernels:

---

webgpu-kernels

/

com.microsoft.GroupQueryAttention

Like 0

Follow

WebGPU Kernels Community 42

{% macro flat_index_2d(workgroupSize, name="i", bound="params.count", guardInline=false) %}
{% set wgTerm = workgroupSize ~ "u" if workgroupSize is number else workgroupSize %}
  // 2D-folded flat index: gid.y carries the high bits past the dispatch's
  // per-axis workgroup fold width.
  let {{ name }} = gid.x + gid.y * {{ DISPATCH_FOLD_WIDTH }}u * {{ wgTerm }};{% endmacro %}
{{ env.wgsl.resourceDeclarations }}

// One invocation per (batch, query head, query token) applies optional per-head
// RMS normalization followed by NeoX half-split rotary embedding at the query's
// absolute position. Present-cache keys are already transformed, so this pass
// writes only Q in BSH layout. f16 queries are widened for normalization,
// trigonometry, and rotation. The result narrows only when the following
// attention pass consumes f16 operand tiles; other configurations store f32.
const HEAD_DIM: u32 = {{ headDim }}u;
const HEAD_DIM_V4: u32 = {{ headDim }}u / 4u;
const Q_HEADS: u32 = {{ qHeads }}u;
const Q_HIDDEN: u32 = {{ qHidden }}u;
const Q_HIDDEN_V4: u32 = {{ qHidden }}u / 4u;
const WG: u32 = {{ copyWorkgroupSize }}u;
{% set HAS_QNORM = hasQNorm is not defined or hasQNorm %}
const HALF: u32 = {{ half }}u;
{% if HAS_QNORM %}
const QK_EPS: f32 = {{ qkEps }};
{% endif %}

@compute @workgroup_size(WG)
fn main(@builtin(global_invocation_id) gid: vec3<u32>) {
{{ flat_index_2d("WG", "qi", "") }}
  let total = params.batch * Q_HEADS * params.qSeq;
  if (qi >= total) { return; }
  let s = qi % params.qSeq;
  let tmp = qi / params.qSeq;
  let h = tmp % Q_HEADS;
  let b = tmp / Q_HEADS;
  let absPos = params.totalSeq - params.qSeq + s;
  let base = (b * params.qSeq + s) * Q_HIDDEN + h * HEAD_DIM;

  var q: array<f32, HEAD_DIM>;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = f32(query[base + d]); }
{% if HAS_QNORM %}
  var ss = 0.0;
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { ss = ss + q[d] * q[d]; }
  let invRms = inverseSqrt(ss / f32(HEAD_DIM) + QK_EPS);
  for (var d = 0u; d < HEAD_DIM; d = d + 1u) { q[d] = q[d] * invRms * f32(q_norm_weight[d]); }
{% endif %}
  for (var d = 0u; d < HALF; d = d + 1u) {
    let cs = f32(cos_cache[absPos * HALF + d]);
    let sn = f32(sin_cache[absPos * HALF + d]);
    let x0 = q[d];
    let x1 = q[d + HALF];
    q[d] = x0 * cs - x1 * sn;
    q[d + HALF] = x1 * cs + x0 * sn;
  }

  let base4 = (b * params.qSeq + s) * Q_HIDDEN_V4 + h * HEAD_DIM_V4;
  for (var c = 0u; c < HEAD_DIM_V4; c = c + 1u) {
    qout[base4 + c] = {{ "vec4<f16>(" if (qPrepF16 is defined and qPrepF16) else "" }}vec4<f32>(q[c * 4u], q[c * 4u + 1u], q[c * 4u + 2u], q[c * 4u + 3u]){{ ")" if (qPrepF16 is defined and qPrepF16) else "" }};
  }
}