Replace CPU Function Calls with GPU Kernel Invocations #697
Conversation
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
…for including sampling latency breakdown in response
Member
CharlieFRuan
left a comment
CharlieFRuan
left a comment
There was a problem hiding this comment.
This is great, thank you so much! Only two minor nits. Afterwards let's merge!
examples/get-started-latency-breakdown/src/get_started_latency_breakdown.ts
Outdated
Show resolved
Hide resolved
CharlieFRuan
approved these changes
Sep 13, 2025
Member
CharlieFRuan
left a comment
CharlieFRuan
left a comment
There was a problem hiding this comment.
LGTM, great work! Let's follow up with on-GPU sampling!
atebites-hub
pushed a commit
to atebites-hub/web-llm
that referenced
this pull request
Oct 4, 2025
…-ai#697) 1. Replace CPU function calls for the following tasks with GPU kernel invocations: - Apply logit bias - Apply penalties to logits - Compute softmax with temperature (sampling will be replaced in a future PR) 2. Fixed bug with repetition penalty not being used in generation config - Added repetition penalty to CompletionCreateParamsBase and ChatCompletionRequestBase interfaces - Updated definition in GenerationConfig and added reference in engine.ts 3. Added additional field in CompletionCreateParamsBase and ChatCompletionRequestBase interfaces to enable logging of time taken for individual steps 4. Added sanity checks for individual steps in sampleTokenFromLogits Performance Comparison: Compared performance for "canonical" flows averaged across 20 runs - No logit_bias - No logitProcessor - Applied penalties - With and without logprobs 1. Before PR performance (without logprobs): ~0.064s per output token (~15.63 decode tokens/s) 2. After PR performance (without logprobs): ~0.066s per output token (~15.15 decode tokens/s) 3. Before PR performance (with logprobs): ~0.052s per output token (~19.23 decode tokens/s) 5. After PR performance (without logprobs): ~0.048s per output token (~20.83 decode tokens/s) Additional Notes: - Need to profile performance of sampleTopPFromLogits vs sampleTopPFromProb on CPU to determine why performance with logprobs is better - Application of logit_bias is much faster on GPU than CPU - There are additional overheads outside of the sampleTokenFromLogits function that make the performance improvement less pronounced (the total time spent in sampleTokenFromLogits is ~0.0117s before the PR and ~0.0076s after the PR)
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
2 participants
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.
Performance Comparison: Compared performance for "canonical" flows averaged across 20 runs
Additional Notes: