Repository navigation
[DRAFT][DeepSeek-V4] Emit fp32 directly in linear_bf16_fp32 (drop .float() copy) - #3
Closed
AMD-yanfeiwang wants to merge 1 commit into
Closed
AMD-yanfeiwang wants to merge 1 commit into
AMD-yanfeiwang wants to merge 1 commit into
Conversation
…at() copy The AITER path of linear_bf16_fp32 computed a bf16 GEMM and then cast the result to fp32 with `.float()`. That cast shows up in profiles as a separate `bfloat16tofloat32_copy` kernel launched once per call in the DSV4 indexer kv_score path. Request fp32 output from the GEMM directly so the up-cast is fused into the epilogue, removing the copy and avoiding rounding the result to bf16 first. Pairs with an AITER change that provides a native fp32-output kernel for these bf16xbf16 shapes; without it tgemm falls back internally (no regression in correctness).
2 tasks
Owner
Author
|
Superseded by the upstream PR against sgl-project/sglang: sgl-project#28877 (same branch/commit). Closing this fork-internal PR to avoid duplication. |
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
1 participant
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.
Motivation
In the DSV4 indexer kv_score path,
linear_bf16_fp32(AITER branch) computes a bf16 GEMM and then casts to fp32 with.float():Profiling DeepSeek-V4 (8xMI355, tp=8) shows this
.float()as a separatebfloat16tofloat32_copykernel launched once per call (per layer, per step), e.g. shapes[256, {512,1024,2048}]. It is launched fromcompute_kv_score->linear_bf16_fp32. It is small individually but recurs every decode step, and computing in bf16 then up-casting also loses precision.Modifications
otype=torch.float32) so the up-cast is fused into the epilogue and the separate copy kernel is removed.Dependency
This is the consumer-side half. To actually avoid the copy (rather than relocate it), AITER must expose a native fp32-output kernel for these bf16xbf16 shapes — companion draft: AMD-yanfeiwang/aiter
fix/dsv4-bf16-fp32-native-gemm. Without it,tgemmfalls back internally (correctness unaffected, but the copy/torch-fallback may remain), which is why this is a draft.Status / TODO
bfloat16tofloat32_copyis gone in the kv_score pathCI States
Latest PR Test (Base): ❌ Run #27837016614
Latest PR Test (Extra): ❌ Run #27837016436