llm_bench profiling

Describe the bug

Running llm_bench --profile for the three speculative-decoding benchmark modes
(spec_draft_prefill, spec_draft_proposal, spec_verify) on Qwen3-8B-nvfp4 + EAGLE3
produces two independent failures, both of which are documented in
docs/source/user_guide/performance/performance-benchmarks.md
as the canonical way to collect spec-decode performance numbers:

Bug 1 — --profile suppresses all E2E timing (throughput / TTFT / latency absent)

When --profile is passed the benchmark returns immediately after the layer-profiling loop and
never runs E2E timing. As a result:

• E2E Time (actual performance): X ms is not printed for any spec mode
• Tokens/sec (E2E) (generation throughput) is not printed for spec_verify /
  spec_draft_proposal
• Prefill throughput (tok/s) is not printed for spec_draft_prefill
• TTFT and per-iteration latency are therefore unobservable with --profile

The performance-benchmarks doc explicitly lists --profile as the flag to use for all three
spec-decode component commands, so users following the documented workflow get no actionable
throughput or latency numbers at all.

Impact: the documented benchmarking procedure produces incomplete output; users cannot
collect TTFT, generation throughput, or latency for any spec-decode phase.

Bug 2 — MHA / GEMM layer attribution = 0% for all spec-decode modes

When layer profiling does run (via --profile), the breakdown consistently reports
MHA = 0.0 ms (0%) and GEMM = 0.0 ms (0%) with Kgen+Other absorbing 100% of the
measured time across all three modes — including spec_verify which runs the full 36-layer base
model with FMHA + NVFP4 GEMMs.

Impact: the layer breakdown is useless for diagnosing compute bottlenecks in
speculative-decoding pipelines.

---

Steps/Code to reproduce bug

Build configuration:

cmake .. \
  -DCMAKE_BUILD_TYPE=Release \
  -DTRT_PACKAGE_DIR=/usr \
  -DCMAKE_TOOLCHAIN_FILE=cmake/aarch64_linux_toolchain.cmake \
  -DEMBEDDED_TARGET=jetson-thor
make -j$(nproc)

Runtime command used:

ENGINE=workspace/Qwen3-8B-nvfp4-eagle3/engine/base

# Speculative decoding component timing
./build/examples/llm/llm_bench \
  --engineDir $ENGINE \
  --mode spec_draft_prefill \
  --batchSize 1 --inputLen 370 \
  --warmup 2 --iterations 10 --profile

./build/examples/llm/llm_bench \
  --engineDir $ENGINE \
  --mode spec_draft_proposal \
  --batchSize 1 --draftTreeSize 10 --draftStep 6 --pastKVLen 370 \
  --warmup 2 --iterations 10 --profile

./build/examples/llm/llm_bench \
  --engineDir $ENGINE \
  --mode spec_verify \
  --batchSize 1 --verifyTreeSize 60 --pastKVLen 370 \
  --warmup 2 --iterations 10 --profile

Observed output — spec_draft_prefill (Bug 1 + Bug 2):

[INFO] Bench Config:
  Mode: eagle_draft_prefill
  Batch Size: 1
  Iterations: 10
  Input Len: 370
  Reuse KV Len: 0

=== Results Summary ===
Layer Profiling Time (for breakdown analysis):
  Total:      3.0258 +/- 0.0658 ms
  MHA:        0.0000 +/- 0.0000 ms (0.0%)   ← Bug 2: expected non-zero
  GEMM:       0.0000 +/- 0.0000 ms (0.0%)   ← Bug 2: expected non-zero
  Kgen+Other: 3.0258 +/- 0.0658 ms (100.0%) ← Bug 2: absorbs entire budget

← Bug 1: no "E2E Time" line
← Bug 1: no "Tokens/sec (E2E)" / prefill throughput line

Observed output — spec_draft_proposal (Bug 1 + Bug 2):

[INFO] Bench Config:
  Mode: eagle_draft_proposal
  Batch Size: 1
  Iterations: 10
  Past KV Len: 370
  Draft Tree Size: 10

=== Results Summary ===
Layer Profiling Time (for breakdown analysis):
  Total:      2.0117 +/- 0.0629 ms
  MHA:        0.0000 +/- 0.0000 ms (0.0%)   ← Bug 2: expected non-zero
  GEMM:       0.0000 +/- 0.0000 ms (0.0%)   ← Bug 2: expected non-zero
  Kgen+Other: 2.0117 +/- 0.0629 ms (100.0%) ← Bug 2: absorbs entire budget

← Bug 1: no "E2E Time" line
← Bug 1: no "Tokens/sec (E2E)" / generation throughput line

Observed output — spec_verify (Bug 1 + Bug 2):

[INFO] Bench Config:
  Mode: eagle_verify
  Batch Size: 1
  Iterations: 10
  Past KV Len: 370
  Verify Tree Size: 60

=== Results Summary ===
Layer Profiling Time (for breakdown analysis):
  Total:      29.0505 +/- 0.1185 ms
  MHA:        0.0000 +/- 0.0000 ms (0.0%)    ← Bug 2: expected non-zero (36-layer base model)
  GEMM:       0.0000 +/- 0.0000 ms (0.0%)    ← Bug 2: expected non-zero
  Kgen+Other: 29.0505 +/- 0.1185 ms (100.0%) ← Bug 2: absorbs entire budget

← Bug 1: no "E2E Time" line
← Bug 1: no "Tokens/sec (E2E)" / generation throughput line

Expected behavior

Bug 1 — E2E timing: --profile should not suppress E2E timing. The output should include
all of the following, consistent with vanilla prefill / decode modes:

# spec_draft_prefill
E2E Time (actual performance): X.XXXX ms
InputLen: 370, ReuseKVLen: 0, ContextLen: 370
Tokens/sec (E2E): XXXX.X          ← draft prefill throughput (TTFT proxy)

# spec_draft_proposal
E2E Time (actual performance): X.XXXX ms
DraftTreeSize: 10, PastKVLen: 370, OSL: N
Tokens/sec (E2E): XXXX.X          ← draft proposal throughput

# spec_verify
E2E Time (actual performance): XX.XXXX ms
VerifyTreeSize: 60, PastKVLen: 370, OSL: N
Tokens/sec (E2E): XXX.X           ← generation throughput (most important number)

Bug 2 — Layer breakdown: --profile should attribute time across all three categories
consistently with vanilla prefill / decode modes. For spec_verify (36-layer base model,
FMHA for attention + NVFP4 GEMMs for MLP) the MHA and GEMM shares should be substantial.
spec_draft_prefill and spec_draft_proposal use the 1-layer EAGLE draft head and should
similarly show non-zero MHA and GEMM time.

---

System information (Edge Device)

• Platform: NVIDIA Thor
• Software release: ?
• CPU architecture: aarch64
• GPU compute capability: SM110
• Total device memory: ? (unified memory; system RAM = 122 GiB)
• Build type: Release
• Library versions:
  • TensorRT Edge-LLM version or commit hash: a44a1a8
  • CUDA: 13.2
  • TensorRT: 10.16.2.10
  • C++ compiler: GCC 13.3.0
• CMake options used:
  • CMAKE_TOOLCHAIN_FILE: cmake/aarch64_linux_toolchain.cmake
  • EMBEDDED_TARGET: jetson-thor
  • TRT_PACKAGE_DIR: /usr
• Engine: Qwen3-8B-nvfp4-eagle3 — base model 36 layers / hiddenSize=4096 / 8 KV heads,
  draft (EAGLE3) 1 layer / hiddenSize=4096 / maxDraftTreeSize=60
• Any other details that may help:
  • Bug 1 root cause hint: llm_bench returns early after runLayerProfilingLoop when
    --profile is set (examples/llm/llm_bench.cpp), skipping the entire E2E timing phase
    (runRepeatedE2ETiming / runSequentialE2ETiming). The e2eTimeMsResult guard
    (if (e2eTimeMsResult > 0)) in logResultsSummary therefore never fires for spec modes
    when --profile is used. Vanilla prefill and decode modes are unaffected because their
    E2E path is separate.
  • Bug 2 root cause hint: AttentionPlugin reports
    "FMHA supported for headSize=128, using FMHA for prefill + XQA for decode" (loaded correctly
    for all 36 layers in spec_verify), so FMHA kernels execute but their time is not being
    attributed to the mhaTimeMs bucket in KernelTimes — it falls entirely into
    kgenOtherTimeMs. The same misattribution affects the 1-layer draft head in
    spec_draft_prefill and spec_draft_proposal.

Start in examples/llm/llm_bench.cpp, tracing runLayerProfilingLoop, runRepeatedE2ETiming, runSequentialE2ETiming, and logResultsSummary for the three documented commands. Reproduce the profile runs, then inspect how AttentionPlugin timings reach the KernelTimes buckets. Done means --profile reports E2E timing and throughput, with MHA and GEMM receiving non-zero attribution where those kernels run.