Skip to content

perf(dense-ffn): per-op profiling breakdown for decode forward()#256

Closed
Sujimoshi wants to merge 1 commit into
Avarok-Cybersecurity:mainfrom
Sujimoshi:perf/dense-ffn-fix
Closed

perf(dense-ffn): per-op profiling breakdown for decode forward()#256
Sujimoshi wants to merge 1 commit into
Avarok-Cybersecurity:mainfrom
Sujimoshi:perf/dense-ffn-fix

Conversation

@Sujimoshi

Copy link
Copy Markdown
Contributor

Summary

Second in a planned split of #218 (DFlash speculative decoding) into smaller, independently reviewable PRs (first was #255).

  • Adds a prof! macro to crates/spark-model/src/layers/dense_ffn.rs (same shape as the existing pattern in moe/forward.rs / qwen3_ssm/ssm_forward.rs) and wraps every GPU op in DenseFfnLayer::forward()'s decode path with it.
  • Covers all four dispatch branches: FP8 fast/slow path, BF16 path, and the NVFP4 default path (use_sw variant, split_silu branch, fused SiLU/GeLU branches).

Note on provenance: the original commit in #218 (18fc6549) targeted an older shape of forward(). Since then, #229 (NVFP4 W4A4 MMQ FFN) restructured the function and added the fp8/use_sw/split_silu branches. This PR re-applies the profiling breakdown against the current (post-#229) structure rather than patching mechanically — more dispatch branches means the per-op breakdown is more useful, not less, for isolating decode cost under whichever config is active.

No behavior change — prof! is a no-op unless ctx.profile is set.

Test plan

  • cargo check -p spark-model (ATLAS_SKIP_BUILD=1, CUDA stubbed) — clean
  • cargo fmt --check -p spark-model — clean
  • cargo clippy -p spark-model --tests (ATLAS_SKIP_BUILD=1) — no warnings

🤖 Generated with Claude Code

Adds a prof! macro (same pattern as moe/forward.rs, qwen3_ssm/ssm_forward.rs)
and wraps every GPU op in DenseFfnLayer::forward()'s decode path with it —
FP8 fast/slow path, BF16 path, and the NVFP4 default path (use_sw variant,
split_silu branch, fused SiLU/GeLU branches).

Split out of Avarok-Cybersecurity#218. Re-applied against the post-Avarok-Cybersecurity#229 forward() (which added
the fp8 fast-path, use_sw, and split_silu branches on top of the original
BF16/NVFP4 dispatch this profiling was written for) rather than patched
mechanically — Avarok-Cybersecurity#229 added more dispatch branches, not fewer, making a
per-op breakdown more useful for isolating decode cost under whichever
env-var/checkpoint configuration is active.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
@Sujimoshi

Copy link
Copy Markdown
Contributor Author

Closing — this was purely diagnostic instrumentation (Dense FFN per-op profiling breakdown) split out of #218. The original perf investigation it supported is already done and documented; no active use case for it right now. Cheap to re-add later against whatever the code looks like at that point if a real need comes up.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant