Skip to content

Tune SM100 GEMM scheduling for dW shapes - #1380

Merged
kolehma8 merged 2 commits into
linkedin:mainfrom
justinhh4:justinhh4/sm100-gemm-shape-tuning
Aug 14, 2026
Merged

Tune SM100 GEMM scheduling for dW shapes#1380
kolehma8 merged 2 commits into
linkedin:mainfrom
justinhh4:justinhh4/sm100-gemm-shape-tuning

Conversation

@justinhh4

Copy link
Copy Markdown
Contributor

Summary

  • keep the fixed two-CTA SM100 GEMM design
  • use swizzle-2 for measured large-M, small-K dW shapes
  • select 4 A/B stages for K=256/2048 and 6 stages for K=512/1024
  • retain the existing 6-stage, swizzle-1 configuration for all other shapes

B200 BF16 results

Representative M=128256, N=4096 results (median latency):

K Before After Speedup cuBLAS
256 0.353 ms 0.252 ms 1.40x 0.264 ms
512 0.511 ms 0.394 ms 1.30x 0.396 ms
1024 0.932 ms 0.699 ms 1.33x 0.685 ms
2048 1.779 ms 1.349 ms 1.32x 1.336 ms

Across 16 Llama, Qwen, and DeepSeek dW shapes, the tuned configuration was 1.07-1.40x faster than the previous configuration and stayed within about 4% of cuBLAS, often slightly faster.

Testing

  • python -m pytest test/cutedsl/test_sm100_gemm.py -q (11 passed)
  • ruff check and ruff format --check on changed files

Justin Hu and others added 2 commits August 14, 2026 01:14
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
@kolehma8
kolehma8 added this pull request to the merge queue Aug 14, 2026
Merged via the queue into linkedin:main with commit ce72ded Aug 14, 2026
5 of 7 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants