Skip to content

[GEMM] Fix sparse split-K metadata offsets (issue #3533) - #3625

Open
XFDG wants to merge 1 commit into
NVIDIA:mainfrom
XFDG:fix/sparse-splitk-metadata-offset-3533
Open

XFDG wants to merge 1 commit into
NVIDIA:mainfrom
XFDG:fix/sparse-splitk-metadata-offset-3533

Conversation

@XFDG

@XFDG XFDG commented Sep 13, 2026

Copy link
Copy Markdown

Summary

Correct the packed metadata E column offset for K slices after the first in three sparse GEMM kernel bodies:

  • SparseGemm
  • SparseGemmWithAbsmax
  • SparseGemmWithEpilogueVisitor

Each offset now divides by kElementsPerElementE, matching the metadata extent and GemmSparseUniversal.

Rationale

The A and B slice offsets are expressed in their respective logical element coordinates. E is different: each stored metadata element describes kElementsPerElementE groups of sparse A values, and its iterator extent already divides by both kSparse and kElementsPerElementE.

The old E offset divided only by kSparse. For split-K slices greater than zero, this started metadata iteration beyond the packed E slice. Predication then suppressed valid metadata loads and the slice accumulated against invalid zero-filled metadata.

Validation

Used the non-universal SparseGemm path from example 15 with serial split-K enabled. Built with CUDA 13.1 and -arch=sm_100a and ran on an NVIDIA B200.

Control:

  • split-K=1 passes the complete example reference comparison on current main

Before:

  • split-K=2 fails the same end-to-end reference comparison

After:

  • split-K=2 passes the end-to-end reference comparison
  • all three changed headers compile in the same CUDA translation unit
  • git diff --check passes

This path is distinct from the existing GemmSparseUniversal split-K test coverage.

Fixes #3533

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.

Sparse split-K kernels offset the E operand without the kElementsPerElementE division

1 participant