perf(gemm): route all bf16 shapes through cublasLt on Kunlun, eliminting costly f32 fallback - #1552
Merged
Merged
Conversation
…ting costly f32 fallback Eliminate the f32-cast fallback for large skinny bf16 GEMMs (batch-1 decode with is_transed, n==1, m>2048). cublasLt (xblas) handles every bf16 shape correctly; the previous f32 path re-cast the entire weight matrix on every token, costing 16–200 ms per projection/lm_head call and making 0.6B-class inference ~80× slower than expected. fix(random_sample): eliminate race conditions in top-k/top-1 kernels on Kunlun Add missing cluster-level barriers (`sync_cluster()` / `mfence()`) in `TopOneKernel` and `TopkKernel`. Without them, thread 0 could read stale workspace values while other cores were still writing, causing non-deterministic wrong-token selection that manifested as garbled text during greedy decode and top-k sampling.
wooway777
approved these changes
Sep 8, 2026
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
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.
Eliminate the f32-cast fallback for large skinny bf16 GEMMs (batch-1
decode with is_transed, n==1, m>2048). cublasLt (xblas) handles every
bf16 shape correctly; the previous f32 path re-cast the entire weight
matrix on every token, costing 16–200 ms per projection/lm_head call
and making 0.6B-class inference ~80× slower than expected.
fix(random_sample): eliminate race conditions in top-k/top-1 kernels on Kunlun
Add missing cluster-level barriers (
sync_cluster()/mfence()) inTopOneKernelandTopkKernel. Without them, thread 0 could readstale workspace values while other cores were still writing, causing
non-deterministic wrong-token selection that manifested as garbled text
during greedy decode and top-k sampling.
昆仑上优化前1054秒，优化后10秒以内
优化后
test_perf测试
优化前性能：
优化后性能：
后续应该仍然有优化空间，profiling看到有蛮多同步操作，后续优化再提PR吧
