[PR #460] [Fix] Bound cuBLASLt descriptor lifetime on errors #464

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

📋 Pull Request Information

Original PR: https://github.com/deepseek-ai/DeepGEMM/pull/460
Author: @ZenAlexa
Created: 9/28/2026
Status: 🔄 Open

Base: main ← Head: contrib/cublaslt-stack-descriptors


📝 Commits (5)

  • 0677683 [Fix] Bound cuBLASLt descriptor lifetime on errors
  • 7d53ec4 [Perf] Drop redundant cuBLASLt scalar type update
  • ad60716 Merge public release 26/09/30
  • 99a3343 Merge upstream main into cuBLASLt descriptor fix
  • d8081b2 Merge current main

📊 Changes

1 file changed (+46 additions, -54 deletions)

View changed files

📝 csrc/jit_kernels/impls/smxx_cublaslt.hpp (+46 -54)

📄 Description

Summary

Replaced the five cuBLASLt descriptor allocations in the shared GEMM path with caller-owned opaque descriptors initialized through the native *Init APIs. The layouts live in the calling scope, and the operation descriptor and preference live through heuristic selection and matmul enqueue.

The current Create sequence leaves its corresponding Destroy calls at the end of the function. A heuristic or matmul error throws through the existing error macros and exits before those calls. Stack-owned descriptors give the success and exception paths the same bounded lifetime.

The native matmul initializer also sets the scalar type, so the shared path now omits the duplicate SCALE_TYPE attribute call. The algorithm selection, per-stream workspace, matrix attributes, scaling pointers, and public interfaces retain their existing behavior. NVIDIA TransformerEngine uses the same native descriptor ownership pattern in its grouped GEMM implementation.

Reproduction

A public cublaslt_gemm_nt call with two CUDA FP64 [16, 16] inputs and an FP32 [16, 16] output reaches heuristic selection and returns CUBLAS_STATUS_NOT_SUPPORTED. Repeating this rejected call retains the allocated descriptors on the current default branch.

Measured glibc mallinfo2().uordblks in separate baseline and candidate processes after 128 warmup rejections. Each checkpoint runs garbage collection and malloc_trim(0); inputs remain allocated throughout. The baseline is 78b69000794d0937b47ae3387eff7663410264d1.

20,000 rejected calls Baseline native allocation growth Candidate native allocation growth
Unsupported dtype at cuBLASLt heuristic selection 11,840,144 bytes 144 bytes
Invalid output shape at the public input boundary 256 bytes 256 bytes

The unsupported-dtype candidate measurement was repeated at 7d53ec4; the invalid-output-shape control was measured at 0677683. The baseline grows by approximately 592 bytes per cuBLASLt rejection. The candidate at 7d53ec4 has a stable native allocation total from 1,000 through 20,000 rejected calls. Both versions preserve the rejection and complete a subsequent valid BF16 GEMM with zero maximum absolute error. At ad60716, a fresh candidate process again retained a constant native allocation total from 1,000 through 20,000 rejected calls (147,976,064 bytes), followed by a valid BF16 GEMM with zero maximum absolute error.

Validation

Native validation at ad60716 on October 8 included main 057ca596. That diff retains the release's output dtype, alpha scaling, and accumulation handling. A fresh native build passed all six existing test functions with their complete internal shape matrices and original tolerances (6 passed, 1,523.68 seconds):

python -m pytest -q \
  tests/test_bf16.py::test_cublaslt_gemm \
  tests/test_bf16.py::test_cublaslt_batched_syrk \
  tests/test_bf16.py::test_cublaslt_batched_symm \
  tests/test_einsum.py::test_bhr_hdr_bhd \
  tests/test_einsum.py::test_bhd_hdr_bhr \
  tests/test_einsum.py::test_bhd_bhr_hdr

These exercise matrix orientations, padded and strided layouts, BF16 and FP32 outputs, alpha scaling, and accumulation. At ad60716, a separate two-stream check uses distinct BF16 inputs and FP32 accumulators with non-aliased outputs. It validates finite reference-matching results after 64 interleaved eager rounds and after 64 CUDA graph replays; both stages pass with calc_diff below 1.2e-13 on each stream.

At 0677683, four paired baseline and candidate workloads passed numerical checks in eager execution and CUDA graph replay: padded-weight bhr,hdr->bhd at B=4, H=128, R=512, D=128, batched SYRK and SYMM at batch=16, M=576, K=5120, and an FP32-output BF16 GEMM at M=4, N=64, K=7168. The measurements use the library's default heuristic cache. Host-call wall times varied across rounds.

Hardware: H200 (driver reports NVIDIA L20X), PyTorch-reported compute capability 9.0, driver 570.172.08. The separate driver capability query reports 8.9. PyTorch 2.12.0.dev20260408+cu128, CUDA toolkit 12.9.86, loaded cuBLASLt 12.8.4 (cublasLtGetVersion() == 120804). The baseline and ad60716 native extensions both built successfully. The shared NVFP4 cuBLASLt wrapper compiled with the extension.

A separate same-process host probe compared descriptor initialization with the repeated scalar-type setter and with the initializer alone. Sixteen alternating rounds of 100,000 initializations gave median thread CPU times of 60.25 ns and 47.45 ns for FAST_TF32, and 60.50 ns and 47.29 ns for FP32. Public descriptor attributes matched for both compute types and all four transpose combinations. These numbers measure host descriptor initialization.

Both native builds use the same CPPFLAGS='-include cuda_fp8.h' setting for the header dependency reported in #439. This PR leaves that separate fix with its existing author.

The timing comparison isolates host descriptor initialization; complete API and model-level performance remain unmeasured. SM100 received compile coverage for the shared wrapper; its runtime behavior requires additional hardware.

Current integration

Merged 5081f8513af5 through a normal merge at d8081b2876e1 on October 23. The descriptor implementation is unchanged. Source review covered all six wrappers, epilogue dispatch, stream-owned workspaces and exception exits; the contribution diff check passed. The native numerical, two-stream/graph, failure-recovery and descriptor-cost results above retain their recorded revisions and hardware coverage.

Current CUDA compilation and SM90/SM100 execution remain unrun. Please run the descriptor failure/recovery cases in the project runtime lanes, including the shared NVFP4 wrapper on SM100.


🔄 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/460 **Author:** [@ZenAlexa](https://github.com/ZenAlexa) **Created:** 9/28/2026 **Status:** 🔄 Open **Base:** `main` ← **Head:** `contrib/cublaslt-stack-descriptors` --- ### 📝 Commits (5) - [`0677683`](https://github.com/deepseek-ai/DeepGEMM/commit/0677683ab9451ab8b6f6c48ee74cc6ed43de7934) [Fix] Bound cuBLASLt descriptor lifetime on errors - [`7d53ec4`](https://github.com/deepseek-ai/DeepGEMM/commit/7d53ec4bf68dbde4f1f4d86613c4d8ce21af0bea) [Perf] Drop redundant cuBLASLt scalar type update - [`ad60716`](https://github.com/deepseek-ai/DeepGEMM/commit/ad6071602a5868df422f153cab10a86e9f9c0f4e) Merge public release 26/09/30 - [`99a3343`](https://github.com/deepseek-ai/DeepGEMM/commit/99a334336cb19e5242b9f01a2276ae643ca01b8c) Merge upstream main into cuBLASLt descriptor fix - [`d8081b2`](https://github.com/deepseek-ai/DeepGEMM/commit/d8081b2876e1c1408c5fe2a46918c527a1cbd1f4) Merge current main ### 📊 Changes **1 file changed** (+46 additions, -54 deletions) <details> <summary>View changed files</summary> 📝 `csrc/jit_kernels/impls/smxx_cublaslt.hpp` (+46 -54) </details> ### 📄 Description ## Summary Replaced the five cuBLASLt descriptor allocations in the shared GEMM path with caller-owned opaque descriptors initialized through the native `*Init` APIs. The layouts live in the calling scope, and the operation descriptor and preference live through heuristic selection and matmul enqueue. The current `Create` sequence leaves its corresponding `Destroy` calls at the end of the function. A heuristic or matmul error throws through the existing error macros and exits before those calls. Stack-owned descriptors give the success and exception paths the same bounded lifetime. The native matmul initializer also sets the scalar type, so the shared path now omits the duplicate `SCALE_TYPE` attribute call. The algorithm selection, per-stream workspace, matrix attributes, scaling pointers, and public interfaces retain their existing behavior. NVIDIA TransformerEngine uses the same native descriptor ownership pattern in its [grouped GEMM implementation](https://github.com/NVIDIA/TransformerEngine/blob/9c33764bd9fd88cb4984ef2b60e9f0623055d3b2/transformer_engine/common/gemm/cublaslt_grouped_gemm.cu#L1067-L1118). ## Reproduction A public `cublaslt_gemm_nt` call with two CUDA FP64 `[16, 16]` inputs and an FP32 `[16, 16]` output reaches heuristic selection and returns `CUBLAS_STATUS_NOT_SUPPORTED`. Repeating this rejected call retains the allocated descriptors on the current default branch. Measured glibc `mallinfo2().uordblks` in separate baseline and candidate processes after 128 warmup rejections. Each checkpoint runs garbage collection and `malloc_trim(0)`; inputs remain allocated throughout. The baseline is `78b69000794d0937b47ae3387eff7663410264d1`. | 20,000 rejected calls | Baseline native allocation growth | Candidate native allocation growth | | --- | ---: | ---: | | Unsupported dtype at cuBLASLt heuristic selection | 11,840,144 bytes | 144 bytes | | Invalid output shape at the public input boundary | 256 bytes | 256 bytes | The unsupported-dtype candidate measurement was repeated at `7d53ec4`; the invalid-output-shape control was measured at `0677683`. The baseline grows by approximately 592 bytes per cuBLASLt rejection. The candidate at `7d53ec4` has a stable native allocation total from 1,000 through 20,000 rejected calls. Both versions preserve the rejection and complete a subsequent valid BF16 GEMM with zero maximum absolute error. At `ad60716`, a fresh candidate process again retained a constant native allocation total from 1,000 through 20,000 rejected calls (147,976,064 bytes), followed by a valid BF16 GEMM with zero maximum absolute error. ## Validation Native validation at `ad60716` on October 8 included main `057ca596`. That diff retains the release's output dtype, alpha scaling, and accumulation handling. A fresh native build passed all six existing test functions with their complete internal shape matrices and original tolerances (6 passed, 1,523.68 seconds): ```bash python -m pytest -q \ tests/test_bf16.py::test_cublaslt_gemm \ tests/test_bf16.py::test_cublaslt_batched_syrk \ tests/test_bf16.py::test_cublaslt_batched_symm \ tests/test_einsum.py::test_bhr_hdr_bhd \ tests/test_einsum.py::test_bhd_hdr_bhr \ tests/test_einsum.py::test_bhd_bhr_hdr ``` These exercise matrix orientations, padded and strided layouts, BF16 and FP32 outputs, alpha scaling, and accumulation. At `ad60716`, a separate two-stream check uses distinct BF16 inputs and FP32 accumulators with non-aliased outputs. It validates finite reference-matching results after 64 interleaved eager rounds and after 64 CUDA graph replays; both stages pass with `calc_diff` below `1.2e-13` on each stream. At `0677683`, four paired baseline and candidate workloads passed numerical checks in eager execution and CUDA graph replay: padded-weight `bhr,hdr->bhd` at `B=4, H=128, R=512, D=128`, batched SYRK and SYMM at `batch=16, M=576, K=5120`, and an FP32-output BF16 GEMM at `M=4, N=64, K=7168`. The measurements use the library's default heuristic cache. Host-call wall times varied across rounds. Hardware: H200 (driver reports NVIDIA L20X), PyTorch-reported compute capability 9.0, driver `570.172.08`. The separate driver capability query reports 8.9. PyTorch `2.12.0.dev20260408+cu128`, CUDA toolkit `12.9.86`, loaded cuBLASLt `12.8.4` (`cublasLtGetVersion() == 120804`). The baseline and `ad60716` native extensions both built successfully. The shared NVFP4 cuBLASLt wrapper compiled with the extension. A separate same-process host probe compared descriptor initialization with the repeated scalar-type setter and with the initializer alone. Sixteen alternating rounds of 100,000 initializations gave median thread CPU times of 60.25 ns and 47.45 ns for FAST_TF32, and 60.50 ns and 47.29 ns for FP32. Public descriptor attributes matched for both compute types and all four transpose combinations. These numbers measure host descriptor initialization. Both native builds use the same `CPPFLAGS='-include cuda_fp8.h'` setting for the header dependency reported in #439. This PR leaves that separate fix with its existing author. The timing comparison isolates host descriptor initialization; complete API and model-level performance remain unmeasured. SM100 received compile coverage for the shared wrapper; its runtime behavior requires additional hardware. ### Current integration Merged `5081f8513af5` through a normal merge at `d8081b2876e1` on October 23. The descriptor implementation is unchanged. Source review covered all six wrappers, epilogue dispatch, stream-owned workspaces and exception exits; the contribution diff check passed. The native numerical, two-stream/graph, failure-recovery and descriptor-cost results above retain their recorded revisions and hardware coverage. Current CUDA compilation and SM90/SM100 execution remain unrun. Please run the descriptor failure/recovery cases in the project runtime lanes, including the shared NVFP4 wrapper on SM100. --- <sub>🔄 This issue represents a GitHub Pull Request. It cannot be merged through Gitea due to API limitations.</sub>
Sign in to join this conversation.
No labels
pull-request
No milestone
No project
No assignees
1 participant
Notifications
Due date
The due date is invalid or out of range. Please use the format "yyyy-mm-dd".

No due date set.

Dependencies

No dependencies set.

Reference
deepseek-ai/DeepGEMM#464
No description provided.