Avoid zero-initializing JAX grouped GEMM outputs - #592
Conversation
|
Important Review skippedDraft detected. Please check the settings in the CodeRabbit UI or the ⚙️ Run configurationConfiguration used: Path: .coderabbit.yaml Review profile: CHILL Plan: Enterprise Run ID: You can disable this status message by setting the Use the checkbox below for a quick retry:
Comment |
|
Closing in favor of #631, which removes zero init for more than just unfused |
Folds in NVIDIA#592's change, which is being closed in favour of this PR so the whole audit lands as one piece. d, c and the workspace are all written before use, so nothing here needed the fill. Adds the poisoned-allocator test for unfused as well: on fresh device memory the untouched tail reads back as zeros whether or not the fill is present, so a test that does not dirty the allocator first passes either way.
|
Closing in favour of #631, which folds this change in and extends the same audit to the rest of the JAX entry points. The mechanism, the measurement protocol and the framing here are yours — #631 credits this PR as its basis. For the record, I reproduced this independently on a separate B200 before building on it. Two things #631 adds that are worth knowing about: The tail change is not observable in a clean test. On fresh device memory the rows past It does not generalize. I tried removing Happy to hand #631 back or split it if you would rather land this one yourself — no attachment to who carries it. |
* jax: stop zero-filling outputs the kernel already writes Follows #592, which established this for grouped/unfused: passing initialized_outputs makes XLA zero-fill the buffer before every dispatch, and that fill scales with the output, so it is the JAX path's dominant host-visible cost. Extends it to the entry points where the same reasoning holds and, more importantly, establishes which ones it does not hold for. Changed: grouped/glu d, c and the workspace are all written before use dense/swiglu ab12 and c are written in full over their extents grouped/wgrad workspace only; the template keeps its zeros because the kernel writes through wgrad_ptrs and never touches it Deliberately unchanged, having tried and reverted each: dense/amax, dense/srelu, grouped/dglu, grouped/dsrelu, discrete_grouped/swiglu and discrete_grouped/dswiglu. Their comments cite the bridge's leading-dim inference rejecting trailing-unit-dim buffers on pure results, and that is still accurate -- dropping the donation there fails the JAX jit-vs-eager tests, and grouped/dglu returns results differing from torch by 1.125. The zeroing is load-bearing beyond the accumulators. Rows at or past padded_offsets[-1] are now unspecified for grouped/glu, as they already are for the torch wrapper's empty_strided outputs. Documented in the docstring. B200, bf16, 8 experts, grouped/glu jax.jit, median of 3 alternating runs (us): 2048x2048x2048 50.6 -> 36.3 4096x2048x2048 49.1 -> 38.6 8192x4096x2048 122.7 -> 108.4 16384x4096x4096 413.4 -> 385.8 At the largest shape that leaves +1.8 us over the kernel, against +29 before. Adds a partial-coverage test that dirties the allocator before dispatch: on fresh device memory the tail reads back as zeros either way, so a test that does not poison it passes whether or not the fill is there. * jax: stop zero-filling grouped/unfused outputs too Folds in #592's change, which is being closed in favour of this PR so the whole audit lands as one piece. d, c and the workspace are all written before use, so nothing here needed the fill. Adds the poisoned-allocator test for unfused as well: on fresh device memory the untouched tail reads back as zeros whether or not the fill is present, so a test that does not dirty the allocator first passes either way. * test: gate the new JAX tests on CuTeDSL JAX availability Both new tests called skip_unless_sm100() without first checking cutlass.jax.is_available(), unlike every other JAX test in the same two files. On an SM100 machine without the CuTeDSL JAX extensions they would have failed rather than skipped. Caught in review by CodeRabbit.
…#631) * jax: stop zero-filling outputs the kernel already writes Follows NVIDIA#592, which established this for grouped/unfused: passing initialized_outputs makes XLA zero-fill the buffer before every dispatch, and that fill scales with the output, so it is the JAX path's dominant host-visible cost. Extends it to the entry points where the same reasoning holds and, more importantly, establishes which ones it does not hold for. Changed: grouped/glu d, c and the workspace are all written before use dense/swiglu ab12 and c are written in full over their extents grouped/wgrad workspace only; the template keeps its zeros because the kernel writes through wgrad_ptrs and never touches it Deliberately unchanged, having tried and reverted each: dense/amax, dense/srelu, grouped/dglu, grouped/dsrelu, discrete_grouped/swiglu and discrete_grouped/dswiglu. Their comments cite the bridge's leading-dim inference rejecting trailing-unit-dim buffers on pure results, and that is still accurate -- dropping the donation there fails the JAX jit-vs-eager tests, and grouped/dglu returns results differing from torch by 1.125. The zeroing is load-bearing beyond the accumulators. Rows at or past padded_offsets[-1] are now unspecified for grouped/glu, as they already are for the torch wrapper's empty_strided outputs. Documented in the docstring. B200, bf16, 8 experts, grouped/glu jax.jit, median of 3 alternating runs (us): 2048x2048x2048 50.6 -> 36.3 4096x2048x2048 49.1 -> 38.6 8192x4096x2048 122.7 -> 108.4 16384x4096x4096 413.4 -> 385.8 At the largest shape that leaves +1.8 us over the kernel, against +29 before. Adds a partial-coverage test that dirties the allocator before dispatch: on fresh device memory the tail reads back as zeros either way, so a test that does not poison it passes whether or not the fill is there. * jax: stop zero-filling grouped/unfused outputs too Folds in NVIDIA#592's change, which is being closed in favour of this PR so the whole audit lands as one piece. d, c and the workspace are all written before use, so nothing here needed the fill. Adds the poisoned-allocator test for unfused as well: on fresh device memory the untouched tail reads back as zeros whether or not the fill is present, so a test that does not dirty the allocator first passes either way. * test: gate the new JAX tests on CuTeDSL JAX availability Both new tests called skip_unless_sm100() without first checking cutlass.jax.is_available(), unlike every other JAX test in the same two files. On an SM100 machine without the CuTeDSL JAX extensions they would have failed rather than skipped. Caught in review by CodeRabbit.
2 participants
Before submitting
pre-commit runand committed any formatting changes.cat-*, one or moremod-*, and oneorig-*.Affected area
Python API or bindings
Summary
Stop zero-initializing the D, C, and workspace outputs of
grouped_gemm_jax_sm100. The output specs and CuTeDSL kernel are unchanged. Document that rows at or beyond the final padded offset are unspecified.Why
Passing
initialized_outputsmakes JAX zero-fill three output buffers before every grouped GEMM dispatch. The grouped GEMM writes the addressed D/C rows, and its helper writes the workspace. Requiring zero-filled padding adds large device operations to the JAX path without benefiting full-coverage calls, while the equivalent PyTorch API allocates outputs withempty_strided.Same-node A/B results on an isolated ComputeLab B200 using
benchmark/grouped_gemm_dispatch/bench_dispatch.pyfromMR 2332
at head
92b46b39099cec904d62aec0de074b8c0c712201:Environment: one B200 (SM100), cuDNN 9.24.0, CUDA 13.3, JAX 0.11.0, and NVIDIA CUTLASS DSL 4.6.2. Each value is the median of five complete harness invocations, alternating parent/head run order. Each invocation uses five warmups followed by 30 calls and one final
block_until_ready(). Median kernel and Torch-execute controls differed by at most 1.5% between variants.Related issues
None.
API and compatibility impact
Rows before
padded_offsets[-1]are unchanged. Rows at or beyond that offset are now explicitly unspecified instead of guaranteed to be zero. No kernel or output-layout changes are included.Testing
uvx pre-commit run --files python/cudnn/gemm/cutedsl/grouped/unfused/jax_api.py(passed)cd test/python && timeout 120 python -m pytest -q -m L0 fe_api/grouped_gemm/test_grouped_gemm_jax.py::test_grouped_gemm_jax_jit_matches_eageron ComputeLab B200 (1 passed)timeout 900 python benchmark/grouped_gemm_dispatch/bench_dispatch.py --api unfused --json <output>on the same ComputeLab B200