perf(dense-ffn): per-op profiling breakdown for decode forward()#256
Closed
Sujimoshi wants to merge 1 commit into
Closed
perf(dense-ffn): per-op profiling breakdown for decode forward()#256Sujimoshi wants to merge 1 commit into
Sujimoshi wants to merge 1 commit into
Conversation
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
requested review from
AzeezIsh,
SeedSource,
rsafier and
tbraun96
as code owners
July 5, 2026 00:46
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. |
4 tasks
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Summary
Second in a planned split of #218 (DFlash speculative decoding) into smaller, independently reviewable PRs (first was #255).
prof!macro tocrates/spark-model/src/layers/dense_ffn.rs(same shape as the existing pattern inmoe/forward.rs/qwen3_ssm/ssm_forward.rs) and wraps every GPU op inDenseFfnLayer::forward()'s decode path with it.use_swvariant,split_silubranch, fused SiLU/GeLU branches).Note on provenance: the original commit in #218 (
18fc6549) targeted an older shape offorward(). Since then, #229 (NVFP4 W4A4 MMQ FFN) restructured the function and added the fp8/use_sw/split_silubranches. 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 unlessctx.profileis set.Test plan
cargo check -p spark-model(ATLAS_SKIP_BUILD=1, CUDA stubbed) — cleancargo fmt --check -p spark-model— cleancargo clippy -p spark-model --tests(ATLAS_SKIP_BUILD=1) — no warnings🤖 Generated with Claude Code