[PR #475] Fix SM120 split-K alignment after partition caps #479

New issue

Open

opened 2026-10-26 21:18:21 +00:00 by bsakaguchi · 0 comments

bsakaguchi commented

2026-10-26 21:18:21 +00:00

Owner

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/475
Author: @Sunt-ing
Created: 10/15/2026
Status: 🔄 Open

Base: nv_dev ← Head: fix/sm120-split-k-sf-alignment

📝 Commits (1)

26f1584 Fix SM120 split-K alignment after partition caps

📊 Changes

2 files changed (+15 additions, -3 deletions)

View changed files

📝 csrc/jit_kernels/heuristics/sm120.hpp (+4 -3)
📝 tests/test_sm120_gemm.py (+11 -0)

📄 Description

The SM120 split-K heuristic checks scale alignment before capping the partition count. Those caps can break alignment and produce incorrect FP8 GEMM results. Move the alignment search after both caps and add regression tests.

Adapts the fix from vllm-project/DeepGEMM#4 to nv_dev, reusing the existing search.

Reproduction and validation

For (M, N, K) = (32, 64, 2560), the heuristic first selects 5 aligned partitions, then caps this to 2. Each partition spans 10 K blocks, violating the kernel's four-block scale-tile alignment.

The regression covers float and packed scales at K=2048/2560/3584/4096. On an RTX 5090, the original and rebuilt revert fail 4 of 8 cases; the fix passes all 8. The default scaling path also fails in a BF16→FP8 linear→residual→SiLU operator chain at width 2560 for both tested seeds, and passes after the fix. Whole-model accuracy impact and production frequency were not measured.

CUDA Graph benchmarks used explicit recipe=(1, 1, 128), per-token scales for both operands, and FP32 output. Median fixed/original latency ratios across six paired rounds show approximately 16% and 44% longer execution for 32×64×2560 and 32×64×3584, and 7% for 32×2560×2560. The original outputs at these points are incorrect. The paired median changes for four correctly computed controls were within 0.2%; 64×2560×2560 improved by approximately 10%.

Base: nv_dev at b64107f2b9599ca76445b7f62eedf66bae1d095b. Candidate: 26f15840f26684c72e1033d2b678b159c7898543.
Environment: RTX 5090, driver 580.105.08, PyTorch 2.13.0+cu130. Build DeepGEMM from source with its CUTLASS and DeepJIT submodules; the build environment requires libdw-dev.

The following public-API example uses default scaling. Run the same script with the original and fixed packages installed:

import torch
import deep_gemm
from deep_gemm.testing import calc_diff
from deep_gemm.utils import per_token_cast_to_fp8, per_block_cast_to_fp8

torch.backends.cuda.matmul.allow_tf32 = False
generator = torch.Generator(device='cuda').manual_seed(17)
width = 2560
x = torch.randn((32, width), generator=generator, device='cuda', dtype=torch.bfloat16)
w = torch.randn((width, width), generator=generator, device='cuda', dtype=torch.bfloat16) / width ** 0.5
residual = torch.randn((32, width), generator=generator, device='cuda', dtype=torch.bfloat16)
a, sa = per_token_cast_to_fp8(x, use_ue8m0=True)
b, sb = per_block_cast_to_fp8(w, use_ue8m0=True)
a_ref = a.float() * sa.repeat_interleave(128, 1)
b_ref = b.float() * sb.repeat_interleave(128, 0).repeat_interleave(128, 1)
reference = a_ref @ b_ref.T
output = torch.empty_like(reference)
deep_gemm.fp8_gemm_nt((a, sa), (b, sb), output)
actual = torch.nn.functional.silu(output + residual.float()).to(torch.bfloat16)
expected = torch.nn.functional.silu(reference + residual.float()).to(torch.bfloat16)
print('GEMM diff:', calc_diff(output, reference))
print('Final diff:', calc_diff(actual, expected))
torch.testing.assert_close(output, reference, rtol=2e-4, atol=1e-4)

Observed GEMM diff: original/rebuilt revert 0.017614289757198476; fixed 1.0103029524088925e-14. The reference decodes the same FP8 operands, separating the computation error from quantization error.

Regression command on the candidate:

PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q tests/test_sm120_gemm.py -k dense_split_k_sf_alignment

Additional validation: two existing Graph/PDL/tail tests passed; 56 existing split-K alpha, accumulation, padded-output and swapped-stride cases also passed on driver 590.48.01; the operator chain covered two seeds and widths 2048/2560/3584/4096 with both default and explicit scaling. Performance used six alternating original/fixed rounds over eight shapes, seven Graph timing samples per point, 200 GEMM calls per sample.

_{🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.}

## 📋 Pull Request Information **Original PR:** https://github.com/deepseek-ai/DeepGEMM/pull/475 **Author:** [@Sunt-ing](https://github.com/Sunt-ing) **Created:** 10/15/2026 **Status:** 🔄 Open **Base:** `nv_dev` ← **Head:** `fix/sm120-split-k-sf-alignment` --- ### 📝 Commits (1) - [`26f1584`](https://github.com/deepseek-ai/DeepGEMM/commit/26f15840f26684c72e1033d2b678b159c7898543) Fix SM120 split-K alignment after partition caps ### 📊 Changes **2 files changed** (+15 additions, -3 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/heuristics/sm120.hpp` (+4 -3) 📝 `tests/test_sm120_gemm.py` (+11 -0) </details> ### 📄 Description The SM120 split-K heuristic checks scale alignment before capping the partition count. Those caps can break alignment and produce incorrect FP8 GEMM results. Move the alignment search after both caps and add regression tests. Adapts the fix from [vllm-project/DeepGEMM#4](https://github.com/vllm-project/DeepGEMM/pull/4) to `nv_dev`, reusing the existing search. <details> <summary>Reproduction and validation</summary> For `(M, N, K) = (32, 64, 2560)`, the heuristic first selects 5 aligned partitions, then caps this to 2. Each partition spans 10 K blocks, violating the kernel's four-block scale-tile alignment. The regression covers float and packed scales at K=2048/2560/3584/4096. On an RTX 5090, the original and rebuilt revert fail 4 of 8 cases; the fix passes all 8. The default scaling path also fails in a BF16→FP8 linear→residual→SiLU operator chain at width 2560 for both tested seeds, and passes after the fix. Whole-model accuracy impact and production frequency were not measured. CUDA Graph benchmarks used explicit `recipe=(1, 1, 128)`, per-token scales for both operands, and FP32 output. Median fixed/original latency ratios across six paired rounds show approximately 16% and 44% longer execution for 32×64×2560 and 32×64×3584, and 7% for 32×2560×2560. The original outputs at these points are incorrect. The paired median changes for four correctly computed controls were within 0.2%; 64×2560×2560 improved by approximately 10%. Base: `nv_dev` at `b64107f2b9599ca76445b7f62eedf66bae1d095b`. Candidate: `26f15840f26684c72e1033d2b678b159c7898543`. Environment: RTX 5090, driver 580.105.08, PyTorch 2.13.0+cu130. Build DeepGEMM from source with its CUTLASS and DeepJIT submodules; the build environment requires `libdw-dev`. The following public-API example uses default scaling. Run the same script with the original and fixed packages installed: ```python import torch import deep_gemm from deep_gemm.testing import calc_diff from deep_gemm.utils import per_token_cast_to_fp8, per_block_cast_to_fp8 torch.backends.cuda.matmul.allow_tf32 = False generator = torch.Generator(device='cuda').manual_seed(17) width = 2560 x = torch.randn((32, width), generator=generator, device='cuda', dtype=torch.bfloat16) w = torch.randn((width, width), generator=generator, device='cuda', dtype=torch.bfloat16) / width ** 0.5 residual = torch.randn((32, width), generator=generator, device='cuda', dtype=torch.bfloat16) a, sa = per_token_cast_to_fp8(x, use_ue8m0=True) b, sb = per_block_cast_to_fp8(w, use_ue8m0=True) a_ref = a.float() * sa.repeat_interleave(128, 1) b_ref = b.float() * sb.repeat_interleave(128, 0).repeat_interleave(128, 1) reference = a_ref @ b_ref.T output = torch.empty_like(reference) deep_gemm.fp8_gemm_nt((a, sa), (b, sb), output) actual = torch.nn.functional.silu(output + residual.float()).to(torch.bfloat16) expected = torch.nn.functional.silu(reference + residual.float()).to(torch.bfloat16) print('GEMM diff:', calc_diff(output, reference)) print('Final diff:', calc_diff(actual, expected)) torch.testing.assert_close(output, reference, rtol=2e-4, atol=1e-4) ``` Observed GEMM diff: original/rebuilt revert `0.017614289757198476`; fixed `1.0103029524088925e-14`. The reference decodes the same FP8 operands, separating the computation error from quantization error. Regression command on the candidate: ```bash PYTEST_DISABLE_PLUGIN_AUTOLOAD=1 python -m pytest -q tests/test_sm120_gemm.py -k dense_split_k_sf_alignment ``` Additional validation: two existing Graph/PDL/tail tests passed; 56 existing split-K alpha, accumulation, padded-output and swapped-stride cases also passed on driver 590.48.01; the operator chain covered two seeds and widths 2048/2560/3584/4096 with both default and explicit scaling. Performance used six alternating original/fixed rounds over eight shapes, seven Graph timing samples per point, 200 GEMM calls per sample. </details> --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>

bsakaguchi added the

pull-request

label

2026-10-26 21:18:21 +00:00

bsakaguchi referenced this issue

2026-10-26 21:18:22 +00:00

[PR #480] Use the allocated TMEM base in the SM100 HC prenorm GEMM #482