Skip to content

Avoid zero-initializing JAX grouped GEMM outputs #592

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Closed
mgoldfarb-nvidia wants to merge 1 commit into from
Closed

Avoid zero-initializing JAX grouped GEMM outputs #592

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 3 additions & 8 deletions python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -28,7 +28,7 @@

from cudnn.datatypes import _convert_to_cutlass_data_type
from cudnn.tensor_adapter import framework_dtype
from cudnn.jax import TensorSpec, call, gemm_operand_spec, zeros_init
from cudnn.jax import TensorSpec, call, gemm_operand_spec
from ..moe_utils import MoEWeightMode
from .moe_grouped_gemm import MoEGroupedGemmBf16Kernel

Expand Down Expand Up @@ -108,9 +108,8 @@ def grouped_gemm_jax_sm100(
row offsets, ``alpha (experts,)`` float32, ``prob (m, 1, 1)`` float32, and
``b_ptrs`` holding per-expert ``(n, k)`` k-major bfloat16 weight base addresses
(packed little-endian uint8, 8 bytes per pointer — or int64 with x64 mode).
Returns ``(d_tensor, c_tensor)`` with ``c_tensor`` None unless ``generate_c``;
rows at/past ``padded_offsets[-1]`` come back zero-filled (the outputs are
donated zero-initialized buffers).
Returns ``(d_tensor, c_tensor)`` with ``c_tensor`` None unless ``generate_c``.
Rows at/past ``padded_offsets[-1]`` are unspecified.
"""
c_dtype = _convert_to_cutlass_data_type(c_dtype)
d_dtype = _convert_to_cutlass_data_type(d_dtype)
Expand Down Expand Up @@ -207,10 +206,6 @@ def grouped_gemm_jax_sm100(
),
input_spec=(operand, None, None, None, _prob_spec()),
output_spec=(operand, operand, None),
# All three donated: c/d for the trailing-unit-dim layout spec (and defined
# bytes past the last offset); the workspace because the helper kernel writes
# the per-expert TMA descriptors into it (XLA inputs are immutable).
initialized_outputs={0: zeros_init, 1: zeros_init, 2: zeros_init},
kernel=kernel,
n=int(n),
k=int(k),
Expand Down