1. Add --profile flag to attention_benchmark to output Highway profiler

BUILD.bazel

@@ -1156,6 +1156,7 @@ cc_binary(
         ":tokenizer",
         "//io",
         "@highway//:abort_header_only",
+        "@highway//:profiler",
         "@highway//:timer",
     ],
 )

deepseek/deepseek.cc

@@ -918,18 +918,24 @@ static HWY_NOINLINE void DeepSeekAttention(size_t num_tokens, size_t layer_idx,
   const size_t num_interleaved = num_tokens * qbatch.Size();
 
   // ---- Projections (tiled MatMul).
-  if (lc.q_lora_rank > 0) {
-    activations.mla_q_a.OverrideCols(lc.q_lora_rank);
-    CallMatMul(att.pre_att_rms_out, layer.mla_q_a, /*add=*/nullptr, env,
-               activations.mla_q_a);
-    RMSNormInplaceBatched</*kPlainWeight=*/true>(layer.mla_q_a_norm,
-                                                 activations.mla_q_a, env.ctx);
-    CallMatMul(activations.mla_q_a, layer.mla_q_b, /*add=*/nullptr, env, att.q);
-  } else {
-    CallMatMul(att.pre_att_rms_out, layer.mla_q_b, /*add=*/nullptr, env, att.q);
+  {
+    GCPP_ZONE(env.ctx, hwy::Profiler::GlobalIdx(),
+              Zones::kGenAttentionComputeQKV);
+    if (lc.q_lora_rank > 0) {
+      activations.mla_q_a.OverrideCols(lc.q_lora_rank);
+      CallMatMul(att.pre_att_rms_out, layer.mla_q_a, /*add=*/nullptr, env,
+                 activations.mla_q_a);
+      RMSNormInplaceBatched</*kPlainWeight=*/true>(
+          layer.mla_q_a_norm, activations.mla_q_a, env.ctx);
+      CallMatMul(activations.mla_q_a, layer.mla_q_b, /*add=*/nullptr, env,
+                 att.q);
+    } else {
+      CallMatMul(att.pre_att_rms_out, layer.mla_q_b, /*add=*/nullptr, env,
+                 att.q);
+    }
+    CallMatMul(att.pre_att_rms_out, layer.mla_kv_a, /*add=*/nullptr, env,
+               activations.mla_kv_a);
   }
-  CallMatMul(att.pre_att_rms_out, layer.mla_kv_a, /*add=*/nullptr, env,
-             activations.mla_kv_a);
 
   const bool is_mtp = (layer_idx >= config.num_layers);
   const size_t start_pos = qbatch.Pos(0);
@@ -995,6 +1001,8 @@ static HWY_NOINLINE void DeepSeekAttention(size_t num_tokens, size_t layer_idx,
   ParallelFor(
       Parallelism::kFlat, num_interleaved * heads, env.ctx, /*cluster_idx=*/0,
       Callers::kAttComputeQKV, [&](size_t task, size_t worker) HWY_ATTR {
+        GCPP_ZONE(env.ctx, worker,
+                  Zones::kGenAttentionDotSoftmaxWeightedSumPar);
         namespace hn = hwy::HWY_NAMESPACE;
         const hn::ScalableTag<float> df;
         const size_t head = task % heads;
@@ -1076,6 +1084,8 @@ static HWY_NOINLINE void DeepSeekAttention(size_t num_tokens, size_t layer_idx,
 
   // ---- Grouped low-rank output projection.
   {
+    GCPP_ZONE(env.ctx, hwy::Profiler::GlobalIdx(),
+              Zones::kGenAttentionSumHeads);
     namespace hn = hwy::HWY_NAMESPACE;
     att.att_out.OverrideCols(heads * qkv_dim);
     const size_t o_groups = lc.o_groups;
@@ -1455,6 +1465,7 @@ struct DeepSeekMoE {
 static HWY_NOINLINE void DeepSeekDenseFFW(const LayerWeightsPtrs& layer,
                                           Activations& activations,
                                           MatMulEnv& env) {
+  GCPP_ZONE(env.ctx, hwy::Profiler::GlobalIdx(), Zones::kGenFFW);
   const LayerConfig& lc = layer.layer_config;
   if (lc.swiglu_limit == 0.0f) {
     FFWNoVit(layer, activations, env);

deepseek/deepseek_spec.cc

@@ -121,6 +121,7 @@ void GenerateDSparkV4(const ModelConfig& config,
 
   // Parallel DSpark Drafter: generates up to max_drafts speculative tokens in a single parallel pass.
   const auto generate_drafts = [&]() HWY_ATTR {
+    PROFILER_ZONE("Gen.MTP.Draft");
     for (size_t j = 0; j < 16; ++j) {
       drafts[j] = 0;
       confidences[j] = 1.0f;
@@ -159,10 +160,13 @@ void GenerateDSparkV4(const ModelConfig& config,
     }
     DeepSeekMaybeInitHCStreams(activations, env);
     activations.ds_snapshot_after = 0;
-    for (size_t layer_idx = 0; layer_idx < weights.c_layers.size();
-         ++layer_idx) {
-      TransformerLayer(block_size, layer_idx, *weights.GetLayer(layer_idx),
-                       activations, qbatch, env);
+    {
+      PROFILER_ZONE("Gen.MTP.VerifyPass");
+      for (size_t layer_idx = 0; layer_idx < weights.c_layers.size();
+           ++layer_idx) {
+        TransformerLayer(block_size, layer_idx, *weights.GetLayer(layer_idx),
+                         activations, qbatch, env);
+      }
     }
     activations.ds_snapshot_after = -1;
     DeepSeekMaybeFinalizeHCStreams(weights, activations, env);
@@ -182,9 +186,11 @@ void GenerateDSparkV4(const ModelConfig& config,
     while (more && num_acc < actual_drafts &&
            next_committed == drafts[num_acc]) {
       ++num_acc;
-      next_committed = Top1OfSoftmax(activations.logits.RowSpan(num_acc)).token;
+      next_committed =
+          Top1OfSoftmax(activations.logits.RowSpan(num_acc)).token;
       if (debug_log) {
-        fprintf(stderr, ", acc draft[%zu]=%d -> top%zu=%d", num_acc - 1, drafts[num_acc - 1], num_acc, next_committed);
+        fprintf(stderr, ", acc draft[%zu]=%d -> top%zu=%d", num_acc - 1,
+                drafts[num_acc - 1], num_acc, next_committed);
       }
       more = emit(next_committed);
     }
@@ -217,6 +223,7 @@ void GenerateDSparkV4(const ModelConfig& config,
 
     if (more) {
       if (num_acc > 0) {
+        PROFILER_ZONE("Gen.MTP.CommitDSparkKV");
         DeepSeekCommitDSparkKV(num_acc, pos, weights, activations, qbatch, env);
         hwy::CopyBytes(activations.dspark_main_hiddens.Row(num_acc),
                        activations.dspark_main_hiddens.Row(0),

deepseek/run_dsv4.cc

@@ -38,6 +38,7 @@
 #include "gemma/kv_cache.h"
 #include "ops/matmul.h"
 #include "util/threading_context.h"
+#include "hwy/profiler.h"
 
 namespace gcpp {
 
@@ -211,6 +212,9 @@ int Main(int argc, char** argv) {
   if (!pending.empty()) fwrite(pending.data(), 1, pending.size(), stdout);
   if (tokenizer) printf("\n");
   fflush(stdout);
+#if PROFILER_ENABLED
+  env.ctx.profiler.PrintResults();
+#endif
   fprintf(stderr, "Done.\n");
   return 0;
 }

evals/attention_benchmark.cc

@@ -31,6 +31,7 @@
 #include "util/mat.h"
 #include "util/threading_context.h"
 #include "hwy/base.h"
+#include "hwy/profiler.h"
 #include "hwy/timer.h"
 
 namespace gcpp {
@@ -305,6 +306,11 @@ int main(int argc, char** argv) {
                                        : 0.0)
                 << " tok/s)\n";
     }
+#if PROFILER_ENABLED
+    std::cout << "\n=== Profiler Breakdown (Decode) ===\n";
+    env.ctx.profiler.PrintResults();
+    std::cout << "===================================\n";
+#endif
     std::cout << "---------------------------------------------------------\n";
   }