Repository: google/gemma.cpp
File: gemma.cc (C++)

Code before:
    }
  });

  if constexpr (kHeads != kKVHeads) {
    constexpr const size_t q_offset = kHeads * kQKVDim * kModelDim;
    constexpr const size_t k_offset = q_offset + 0 * kQKVDim * kModelDim;
    constexpr const size_t v_offset = q_offset + 1 * kQKVDim * kModelDim;
    const size_t kv_offset = pos * kCachePosSize + layer * kCacheLayerSize;

    ProjKV(k_offset, v_offset, kv_offset);
  }

  pool.Run(0, kHeads, [&](const uint64_t head, size_t /*thread*/) HWY_ATTR {
    // Calculate scores
    float* HWY_RESTRICT q =
        activations.q.data() + head * kQKVDim + batch_idx * kHeads * kQKVDim;
    float* HWY_RESTRICT head_att = activations.att.data() +
                                   head * TConfig::kSeqLen +
                                   batch_idx * kHeads * kQKVDim;

    Rope(q, kQKVDim, pos);
    MulByConst(kQueryScale, q, kQKVDim);
    // Compute Q dot K scores
    for (size_t pos2 = 0; pos2 <= pos; ++pos2) {
      const size_t cache_offset =
          kHeads == kKVHeads
              ? pos2 * kCachePosSize + layer * kCacheLayerSize + head * kQKVDim
              : pos2 * kCachePosSize + layer * kCacheLayerSize;
      const float* HWY_RESTRICT k2 = kv_cache.key_cache.get() + cache_offset;
      const float score = Dot(q, k2, kQKVDim);
      head_att[pos2] = score;
    }
    Softmax(head_att, pos + 1);

    // Weighted summation
    float* HWY_RESTRICT att_out = activations.att_out.data() + head * kQKVDim +
                                  batch_idx * kHeads * kQKVDim;
    hwy::ZeroBytes(att_out, kQKVDim * sizeof(*att_out));
    for (size_t pos2 = 0; pos2 <= pos; ++pos2) {
      const size_t cache_offset =
          kHeads == kKVHeads
              ? pos2 * kCachePosSize + layer * kCacheLayerSize + head * kQKVDim
              : pos2 * kCachePosSize + layer * kCacheLayerSize;
      float* HWY_RESTRICT v2 = kv_cache.value_cache.get() + cache_offset;
      MulByConstAndAdd(head_att[pos2], v2, att_out, kQKVDim);
    }
    // linear projection from kQKVDim back to kModelDim, sum projections
    // across heads
    float* HWY_RESTRICT head_out =
        head == 0
            ? activations.att_post2.data() + batch_idx * kModelDim

Review comment (performance):
This is probably fine as long as we can verify the performance. Might consider turning head * kQKVDim into a const (eg head_offset) outside the for loop, setting it to 0 if kHeads == kKVHeads. That would get rid of conditional branching in the pos2 loop. tagging @jan-wassenberg in case you have other thoughts.

In the future as support is added for other model variations, we'll want to avoid branching in the inference path and have more flexibility for swapping in different inference codepaths for (but for MQA that's probably overkill)

Code after:

    // Compute Q dot K scores
    for (size_t pos2 = 0; pos2 <= pos; ++pos2) {
      const size_t cache_offset =
          pos2 * kCachePosSize + layer * kCacheLayerSize + head_offset;
      const float* HWY_RESTRICT k2 = kv_cache.key_cache.get() + cache_offset;
      const float score = Dot(q, k2, kQKVDim);
      head_att[pos2] = score;
    }
    Softmax(head_att, pos + 1);

    // Weighted summation
    float* HWY_RESTRICT att_out = activations.att_out.data() + head * kQKVDim +
                                  batch_idx * kHeads * kQKVDim;
    hwy::ZeroBytes(att_out, kQKVDim * sizeof(*att_out));
    for (size_t pos2 = 0; pos2 <= pos; ++pos2) {
      const size_t cache_offset =
          pos2 * kCachePosSize + layer * kCacheLayerSize + head_offset;
      float* HWY_RESTRICT v2 = kv_cache.value_cache.get() + cache_offset;
      MulByConstAndAdd(head_att[pos2], v2, att_out, kQKVDim);
    }
    // linear projection from kQKVDim back to kModelDim, sum projections
    // across heads
    float* HWY_RESTRICT head_out =
        head == 0
            ? activations.att_post2.data() + batch_idx * kModelDim
            : activations.att_post1.data() + head * kBatchSize * kModelDim;
    MatVecLoop<kModelDim, kQKVDim>(c_layer->c_attn_vec_einsum_w,
                                   head * kModelDim * kQKVDim, att_out,
                                   head_out);
  };

  if constexpr (kHeads == kKVHeads) {
    // Multi-Head Attention
    pool.Run(0, kHeads, [&](const uint64_t head, size_t /*thread*/) HWY_ATTR {
      const size_t head_offset = head * 3 * kQKVDim * kModelDim;

      ProjQ(head, head_offset);

      const size_t k_offset = head_offset + 1 * kQKVDim * kModelDim;
      const size_t v_offset = head_offset + 2 * kQKVDim * kModelDim;
      const size_t kv_offset =
          pos * kCachePosSize + layer * kCacheLayerSize + head * kQKVDim;

      ProjKV(k_offset, v_offset, kv_offset);

      Attn(head, head * kQKVDim);
    });
  } else {
    // Multi-Query Attention
    pool.Run(0, kHeads, [&](const uint64_t head, size_t /*thread*/) HWY_ATTR {