From 61cc7d085c9c3454e6d47beb5dc7eae68e953903 Mon Sep 17 00:00:00 2001 From: Krzysztof Rymski Date: Fri, 4 Sep 2026 09:44:11 -0700 Subject: [PATCH] 1. Add --profile flag to attention_benchmark to output Highway profiler zone breakdowns when compiled with -DPROFILER_ENABLED=1. 2. Link //third_party/highway:profiler into attention_benchmark in BUILD. 3. Instrument DSpark speculative decoding stages (Gen.MTP.Draft, Gen.MTP.VerifyPass, Gen.MTP.CommitDSparkKV) in deepseek_spec.cc with PROFILER_ZONE markers. 4. Add GCPP_ZONE profiling instrumentation to DeepSeek V4 MLA attention (ComputeQKV, DotSoftmax, SumHeads) and dense FFW in deepseek.cc. 5. Add PrintResults() call under #if PROFILER_ENABLED to run_dsv4.cc. PiperOrigin-RevId: 976362424 --- BUILD.bazel | 1 + deepseek/deepseek.cc | 33 ++++++++++++++++++++++----------- deepseek/deepseek_spec.cc | 19 +++++++++++++------ deepseek/run_dsv4.cc | 4 ++++ evals/attention_benchmark.cc | 6 ++++++ 5 files changed, 46 insertions(+), 17 deletions(-) diff --git a/BUILD.bazel b/BUILD.bazel index 457400f6..60bc61d6 100644 --- a/BUILD.bazel +++ b/BUILD.bazel @@ -1156,6 +1156,7 @@ cc_binary( ":tokenizer", "//io", "@highway//:abort_header_only", + "@highway//:profiler", "@highway//:timer", ], ) diff --git a/deepseek/deepseek.cc b/deepseek/deepseek.cc index f417c320..6024b3cc 100644 --- a/deepseek/deepseek.cc +++ b/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(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( + 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 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); diff --git a/deepseek/deepseek_spec.cc b/deepseek/deepseek_spec.cc index b5f66860..ae896cbe 100644 --- a/deepseek/deepseek_spec.cc +++ b/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), diff --git a/deepseek/run_dsv4.cc b/deepseek/run_dsv4.cc index c1f06b9d..3ab31fb9 100644 --- a/deepseek/run_dsv4.cc +++ b/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; } diff --git a/evals/attention_benchmark.cc b/evals/attention_benchmark.cc index e87bda20..5893e446 100644 --- a/evals/attention_benchmark.cc +++ b/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"; }