Skip to content

[GEMM] Fix EllGemm column-major split-K workspace sizing (issue #3539) - #3624

Open
XFDG wants to merge 1 commit into
NVIDIA:mainfrom
XFDG:fix/ell-gemm-workspace-shape-3539
Open

XFDG wants to merge 1 commit into
NVIDIA:mainfrom
XFDG:fix/ell-gemm-workspace-shape-3539

Conversation

@XFDG

@XFDG XFDG commented Sep 13, 2026

Copy link
Copy Markdown

Summary

Make the ColumnMajor EllGemm specialization size its serial split-K semaphore workspace from the same transposed problem shape used by initialize().

Add a regression test for both rectangular orientations, covering the previous under-allocation and over-allocation cases.

Rationale

ColumnMajor output is implemented by exchanging M/N and forwarding to a RowMajor underlying operator. initialize() correctly computes its launch grid from (n, m, k), but get_workspace_size() used the original (m, n, k).

For rectangular problems, the two shapes can produce different tile counts. With m=264, n=136, ell_blocksize=32, and four split-K slices, workspace sizing returned 15 semaphore slots while the initialized grid contains 18 output tiles. This left part of the semaphore array outside the allocation/initialization contract.

Using the transposed problem shape in get_workspace happens to match initialize() and the underlying operator for every aspect ratio.

Validation

Compiled the same standalone reproducer before and after this change with CUDA 13.1 and -arch=sm_100a, then ran it on an NVIDIA B200.

Before:

  • initialized launch grid: 2 x 9 tiles
  • get_workspace_size reported 60 bytes instead of 72
  • initialize() returned success but cudaMemsetAsync zeroed only 60 of the required 72 bytes

After:

  • get_workspace_size reports 72 bytes
  • initialize() zeroes all 72 bytes
  • the opposite 136 x 264 aspect ratio is also pinned by the new unit test at 60 bytes
  • the new unit-test translation unit compiles successfully for sm_100a
  • git diff --check passes

Fixes #3539

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.

EllGemm ColumnMajor-output workspace sizing uses the unswapped problem shape and undersizes split-K semaphores

1 participant