---
id: pr-sglang-19880
repo: sgl-project/sglang
pr: 19880
title: '[JIT Kernel][Feature] Support JIT custom all reduce (rewrite as v2)'
author: DarkSharpness
date: '2026-03-20'
url: https://github.com/sgl-project/sglang/pull/19880
source_category: upstream-code
architectures:
- sm90
tags:
- PDL
- benchmark_test
- compiler_runtime
- gemm_quant
- hopper
- memory_primitives
techniques:
- fine-grained-quantization
hardware_features:
- hopper
- pdl
kernel_types:
- compiler
- gemm
- memory
languages:
- cuda-cpp
captured_at: '2026-05-16'
status: merged
merge_sha: 2dd9196079e3424bb44bea2ce7697eeed5578e6a
inclusion_reason: Mine shape sets, tolerance rules, warmup logic, and profile commands before mutating code.
changed_paths:
- python/sglang/jit_kernel/.clang-format
- python/sglang/jit_kernel/all_reduce.py
- python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py
- python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh
- python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh
- python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh
- python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh
- python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh
- python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh
- python/sglang/jit_kernel/include/sgl_kernel/ffi.h
- python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
- python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
- python/sglang/jit_kernel/tests/test_custom_all_reduce.py
- python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
- python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py
- python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
artifact_dir: evidence/pull-bundles/sglang/gh-19880
head_ref_oid: 60a406d0ee77c76c96ff013ea36a5cb7a39788d2
---
# [JIT Kernel][Feature] Support JIT custom all reduce (rewrite as v2)

- Repository: `sgl-project/sglang`
- PR: [19880](https://github.com/sgl-project/sglang/pull/19880)
- Decision: `Mine shape sets, tolerance rules, warmup logic, and profile commands before mutating code.`
- Reason: Mine shape sets, tolerance rules, warmup logic, and profile commands before mutating code.

## Changed Paths

- `python/sglang/jit_kernel/.clang-format`
- `python/sglang/jit_kernel/all_reduce.py`
- `python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py`
- `python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh`
- `python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh`
- `python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh`
- `python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh`
- `python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh`
- `python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh`
- `python/sglang/jit_kernel/include/sgl_kernel/ffi.h`
- `python/sglang/jit_kernel/include/sgl_kernel/utils.cuh`
- `python/sglang/jit_kernel/include/sgl_kernel/vec.cuh`
- `python/sglang/jit_kernel/tests/test_custom_all_reduce.py`
- `python/sglang/srt/distributed/device_communicators/custom_all_reduce.py`
- `python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py`
- `python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py`

## Evidence Bundle

`evidence/pull-bundles/sglang/gh-19880`
