Skip to content

[FA3 Attn Backend] Remove Unnecessary Device Sync for FA3 #4745

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 12 commits into from
Mar 27, 2025
Merged

[FA3 Attn Backend] Remove Unnecessary Device Sync for FA3 #4745

merged 12 commits into from
Mar 27, 2025

Conversation

hebiao064
Copy link
Collaborator

@hebiao064 hebiao064 commented Mar 25, 2025 •
edited
Loading

Motivation

  • Remove Unnecessary Device Sync from this comment
  • We found that when extend_no_prefix is false, we used the seqlens_in_batch.max.item as max_seq_len_q, which is not quite right, we should use extend_no_prefix

Thanks @Fridge003 on helping out by providing profiling tips.

Before: 39us
Screenshot 2025-03-26 at 8 47 46 PM

After: 18us
Screenshot 2025-03-26 at 8 47 30 PM

Benchmark

(venv) jobuser [ ~/sglang ]$ python benchmark/gsm8k/bench_sglang.py --num-shots 8 --num-questions 1319 --parallel 1319
100%|████████████████████████████████████████████████| 1319/1319 [00:19<00:00, 68.15it/s]
Accuracy: 0.793
Invalid: 0.002
Latency: 17.409 s
Output throughput: 6893.539 token/s
python -m sglang.bench_one_batch --model /shared/public/models/Meta-Llama-3-8B-Instruct --batch-size 16 --input 1024 --output 512 --attention-backend fa3

# Disabled Cuda Graph
Prefill. latency: 0.37211 s, throughput:  44029.64 token/s
Decode.  median latency: 0.01244 s, median throughput:   1286.33 token/s
Total. latency:  6.756 s, throughput:   3637.85 token/s

# Enabled Cuda Graph
Prefill. latency: 0.37050 s, throughput:  44221.03 token/s
Decode.  median latency: 0.00814 s, median throughput:   1964.89 token/s
Total. latency:  4.527 s, throughput:   5428.93 token/s

Modifications

Modify seqlens_in_batch.max().item() to be seqlens_in_batch.max()

Checklist

@hebiao064 hebiao064 changed the title Remove Unnecessary Device Sync [FA3 Attn Backend] Remove Unnecessary Device Sync Mar 25, 2025
@zhyncs
Copy link
Member

zhyncs commented Mar 25, 2025

ref #4577

@hebiao064
Copy link
Collaborator Author

ref #4577

interesting, let me do a profiling

@zhyncs
Copy link
Member

zhyncs commented Mar 25, 2025

https://github.com/sgl-project/sglang/blob/main/docs/references/benchmark_and_profiling.md#profile-with-pytorch-profiler

@hebiao064 hebiao064 changed the title [FA3 Attn Backend] Remove Unnecessary Device Sync [Do not merge][FA3 Attn Backend] Remove Unnecessary Device Sync Mar 25, 2025
@hebiao064 hebiao064 marked this pull request as draft March 25, 2025 07:06
@hebiao064 hebiao064 force-pushed the bhe/remove_device_sync branch from c29625c to 160c312 Compare March 25, 2025 07:08
hebiao064
hebiao064 commented Mar 25, 2025
View reviewed changes
hebiao064 and others added 2 commits March 25, 2025 20:54
@hebiao064
Merge branch 'main' into bhe/remove_device_sync
6883b53
@hebiao064 @yubofredwang
add item() back
aa4db97 
Co-authored-by: Yubo Wang <yubowang2019@gmail.com>
@hebiao064 hebiao064 changed the title [Do not merge][FA3 Attn Backend] Remove Unnecessary Device Sync [FA3 Attn Backend] Use extend_seq_lens instead of seqlens_in_batch for prefill Mar 26, 2025
@hebiao064 hebiao064 marked this pull request as ready for review March 26, 2025 03:58
@hebiao064
Copy link
Collaborator Author

@merrymercy @zhyncs after investigation, we decided to keep the item() call, please review. We put the details in this PR's decsription

@yubofredwang
Copy link
Contributor

yubofredwang commented Mar 26, 2025 •
edited
Loading

Adding some context here.

According to profiling, item() is implicitly called whenever a tensor is getting used to slice/index a tensor. That is why there are 4 more Device to Host Copy.

Due to the operations after getting max_seq_len_k, there will be two additional kernel launches to do slicing of the page table.

metadata.max_seq_len_k = seqlens_in_batch.max()
metadata.page_table[:, metadata.max_seq_len_k :].fill_(0)
metadata.page_table[:, : metadata.max_seq_len_k].copy_(
    self.req_to_token[req_pool_indices[:bs], : metadata.max_seq_len_k]
)
profile_result

The operation can be potentially fused into one where operation with a mask. Doing benchmark on that.

yubofredwang
yubofredwang reviewed Mar 26, 2025
View reviewed changes
yubofredwang
yubofredwang reviewed Mar 26, 2025
View reviewed changes
@hebiao064
Copy link
Collaborator Author

Test errors are mostly like this

huggingface_hub.errors.HfHubHTTPError: 429 Client Error: Too Many Requests for url: https://huggingface.co/api/models/meta-llama/Llama-2-7b-chat-hf

not related to my change

@hebiao064 hebiao064 changed the title [FA3 Attn Backend] Use extend_seq_lens instead of seqlens_in_batch for prefill [FA3 Attn Backend] Remove Unnecessary Device Sync for FA3 Mar 27, 2025
@hebiao064 hebiao064 mentioned this pull request Mar 27, 2025
Merged
11 tasks
@zhyncs zhyncs merged commit 1b9175c into sgl-project:main Mar 27, 2025
0 of 19 checks passed
@hebiao064 hebiao064 mentioned this pull request Mar 30, 2025
Closed
15 tasks
jimoosciuc pushed a commit to Furion-cn/sglang that referenced this pull request Apr 17, 2025
@hebiao064 @yubofredwang @jimoosciuc
[FA3 Attn Backend] Remove Unnecessary Device Sync for FA3 (sgl-projec…
8b2d1ba 
…t#4745)

Co-authored-by: Yubo Wang <yubowang2019@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Reviewers

@merrymercy merrymercy Awaiting requested review from merrymercy merrymercy is a code owner

@Ying1123 Ying1123 Awaiting requested review from Ying1123 Ying1123 is a code owner

@zhyncs zhyncs Awaiting requested review from zhyncs zhyncs is a code owner

@ispobock ispobock Awaiting requested review from ispobock ispobock is a code owner

@HaiShaw HaiShaw Awaiting requested review from HaiShaw HaiShaw is a code owner

@hnyls2002 hnyls2002 Awaiting requested review from hnyls2002 hnyls2002 is a code owner

@xiezhq-hermann xiezhq-hermann Awaiting requested review from xiezhq-hermann xiezhq-hermann is a code owner

1 more reviewer

@yubofredwang yubofredwang yubofredwang left review comments

Reviewers whose approvals may not affect merge requirements
Assignees

@merrymercy merrymercy

@ispobock ispobock

Labels
None yet
Projects
None yet
Milestone
No milestone
Development

Successfully merging this pull request may close these issues.
5 participants
@hebiao064 @zhyncs @yubofredwang @merrymercy @ispobock