Add quantized-pair rowwise symmetric-memory kernels#785
Merged
AkshitaB merged 3 commits intoJul 19, 2026
Conversation
Port the direct-NVSHMEM quantized-pair rowwise transport kernels and expose them at parity with the reference implementation: dispatchRowsPutPair, dispatchRowsPutScaledWeighted, and gatherRowsGetPair (plus the MXFP8 quant device helpers they use), their host launchers, pybind entries, and Python wrappers. Also restores the collective-launch preflight blocks for these kernels. These fuse the FP8 quantized values and their block scales into a single symmetric put/get (vs the materialized path that stages q and scales separately). They are a reference/tuning path — compiled, callable, and preflight-validated, but not wired into the rowwise forward, which continues to use the materialized transport. Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
The quantized-pair rowwise kernels reference ROWWISE_MXFP8_BLOCK_SIZE, ROWWISE_MXFP8_Q_CHUNK_BLOCKS, and ROWWISE_MXFP8_MAX_SCALE_GROUPS, which live in the shared symm-mem common header. Add the three defines so the extension builds. Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
The kernel's per-warp shared staging buffers need ~80 KB of static shared memory, over the 48 KB limit on pre-SM100 GPUs, so device linking failed for the sm_90 build. It is an SM100+ reference/tuning kernel that is fail-closed at runtime below SM100, so compile a trap stub for __CUDA_ARCH__ < 1000 and keep the full body for SM100+. Co-Authored-By: Tianhua Tao <tianhuatao@gmail.com> Co-Authored-By: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Adds the direct-NVSHMEM quantized-pair rowwise expert-parallel transport kernels as a reference/tuning path.
What
olmo_symm_mem_rowwise_kernels.cuh):dispatchRowsPutPair,dispatchRowsPutScaledWeighted, andgatherRowsGetPair, plus the MXFP8 quantize device helpers they use.olmo_symm_mem_rowwise.cuh):rowwise_dispatch_put_pair,rowwise_dispatch_put_scaled_weighted,rowwise_gather_get_pair.olmo_symm_mem_bindings.cpp,symm_mem_vdev2d.py): the corresponding pybind entries and@torch.compiler.disablePython wrappers.olmo_symm_mem_runtime.cuh): restores the collective-launch grid validation for these kernels.Why
These fuse the FP8 quantized values and their block scales into a single symmetric
nvshmemput/get (the current production path stagesqandscalesas separate buffers;dispatchRowsPutScaledWeightedadditionally folds top-k weighting + scaling into the dispatch). On tested hardware this direct path is bit-exact against the materialized path but not faster at the usual block-count settings, so it is provided as a reference / tuning target: compiled, callable, and preflight-validated, but not wired into the rowwise forward, which continues to use the materialized transport.Testing
Python side (import +
make checks) is clean. The kernels compile GPU-side and are grid-validated by the symmetric-memory extension build in the MoE-kernels GPU CI job.