diff --git a/python/sglang/jit_kernel/.clang-format b/python/sglang/jit_kernel/.clang-format
index 56acfb8b8..690cc3fea 100644
--- a/python/sglang/jit_kernel/.clang-format
+++ b/python/sglang/jit_kernel/.clang-format
@@ -17,7 +17,7 @@ PenaltyReturnTypeOnItsOwnLine: 100    # Keeps return type with function name
 IncludeCategories:
   - Regex: '^<sgl_kernel/.*\.h>$'
     Priority: 0
-  - Regex: '^<sgl_kernel/impl/.*>$'
+  - Regex: '^<sgl_kernel/.*/.*>$'
     Priority: 2
   - Regex: '^<sgl_kernel/.*\.cuh>$'
     Priority: 1
diff --git a/python/sglang/jit_kernel/all_reduce.py b/python/sglang/jit_kernel/all_reduce.py
new file mode 100644
index 000000000..76eca0581
--- /dev/null
+++ b/python/sglang/jit_kernel/all_reduce.py
@@ -0,0 +1,192 @@
+from __future__ import annotations
+
+import enum
+from typing import TYPE_CHECKING, List, NamedTuple, Optional, Tuple, cast
+
+import torch
+import tvm_ffi
+
+from sglang.jit_kernel.utils import (
+    cache_once,
+    is_arch_support_pdl,
+    load_jit,
+    make_cpp_args,
+)
+
+
+class ConfigResult(NamedTuple):
+    num_blocks: int
+    num_threads: int
+
+
+class AllReduceAlgo(enum.Enum):
+    ONE_SHOT_PUSH = enum.auto()
+    ONE_SHOT_PULL = enum.auto()
+    TWO_SHOT_PULL = enum.auto()
+
+    def is_push(self) -> bool:
+        return self == AllReduceAlgo.ONE_SHOT_PUSH
+
+    @property
+    def shot(self) -> int:
+        return 2 if self == AllReduceAlgo.TWO_SHOT_PULL else 1
+
+
+if TYPE_CHECKING:
+    CUSTOM_AR_HANDLE = List[int]
+    CUSTOM_AR_PAIR = Tuple[int, CUSTOM_AR_HANDLE]
+
+    class CustomAllReduceObj:
+        def __init__(
+            self,
+            rank: int,
+            world_size: int,
+            pull_buffer_bytes: int,
+            push_buffer_bytes: int,
+            graph_input_count: int,
+            *,
+            max_pull_blocks: Optional[int] = None,
+            max_push_blocks: Optional[int] = None,
+        ) -> None:
+            """
+            Create a CustomAllReduceObj instance.
+
+            :param rank: The rank of the current process.
+            :param world_size: The total number of processes in the group.
+            :param pull_buffer_bytes: The size of the buffer (in bytes) used for pull-based all-reduce.
+            :param push_buffer_bytes: The size of the buffer (in bytes) used for push-based all-reduce.
+            :param graph_input_count: The maximum number of inputs in all CUDA graphs.
+            :param max_pull_blocks: The maximum number of thread blocks to launch for pull-based all-reduce.
+                                    If None, it will be determined by the implementation.
+            :param max_push_blocks: The maximum number of thread blocks to launch for push-based all-reduce.
+                                    If None, it will be determined by the implementation.
+            """
+
+        @property
+        def world_size(self) -> int: ...
+        def share_storage(self) -> CUSTOM_AR_HANDLE: ...
+        def share_graph_inputs(self) -> List[CUSTOM_AR_PAIR]: ...
+        def post_init(self, handles: List[CUSTOM_AR_HANDLE]) -> None: ...
+        def register_inputs(self, handles: List[List[CUSTOM_AR_PAIR]]) -> None: ...
+        def set_cuda_graph_capture(self, is_capturing: bool) -> None: ...
+        def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None: ...
+        def all_reduce(
+            self, input: torch.Tensor, algo: AllReduceAlgo
+        ) -> tvm_ffi.Tensor: ...
+        def config_pull(
+            self, num_blocks: int = -1, num_threads: int = -1
+        ) -> ConfigResult:
+            """
+            Configure the CUDA kernel's grid and block dimensions.
+            This provides only the upper bound of the configuration,
+            and the actual launch configuration may be determined by implementation.
+            Note that push-based all-reduce can not be configured currently.
+
+            :param num_blocks: The maximum number of thread blocks to launch. -1 means no limit.
+            :param num_threads: The maximum number of threads per block. -1 means no limit.
+
+            :return: The previous configuration as a ConfigResult named tuple.
+            """
+            ...
+
+
+@cache_once
+def _jit_custom_all_reduce_pull_module(dtype: torch.dtype, world_size: int):
+    args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
+    return load_jit(
+        "custom_all_reduce",
+        *args,
+        extra_ldflags=["-lcuda"],
+        cuda_files=["distributed/custom_all_reduce_pull.cuh"],
+        cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
+    )
+
+
+@cache_once
+def _jit_custom_all_reduce_push_module(dtype: torch.dtype, world_size: int):
+    args = make_cpp_args(dtype, world_size, is_arch_support_pdl())
+    return load_jit(
+        "custom_all_reduce",
+        *args,
+        extra_ldflags=["-lcuda"],
+        cuda_files=["distributed/custom_all_reduce_push.cuh"],
+        cuda_wrappers=[("all_reduce", f"custom_all_reduce<{args}>")],
+    )
+
+
+@cache_once
+def get_custom_all_reduce_cls() -> type[CustomAllReduceObj]:
+    module = load_jit(
+        "custom_all_reduce_base",
+        extra_ldflags=["-lcuda"],
+        cuda_files=["distributed/custom_all_reduce_base.cuh"],
+        cuda_wrappers=[("register_once", "register_custom_all_reduce")],
+    )
+    module.register_once()
+    device = torch.cuda.current_device()
+    props = torch.cuda.get_device_properties(device)
+    NUM_CTA = props.multi_processor_count
+    MAX_THREADS = 512
+
+    @tvm_ffi.register_object("sgl.CustomAllReduce")
+    class CustomAllReduceObjReal(tvm_ffi.Object):
+        def __init__(
+            self,
+            rank: int,
+            world_size: int,
+            pull_buffer_bytes: int,
+            push_buffer_bytes: int,
+            graph_input_count: int,
+            *,
+            max_pull_blocks: Optional[int] = None,
+            max_push_blocks: Optional[int] = None,
+        ) -> None:
+            self.__ffi_init__(
+                rank,
+                world_size,
+                NUM_CTA if max_pull_blocks is None else max_pull_blocks,
+                NUM_CTA if max_push_blocks is None else max_push_blocks,
+                pull_buffer_bytes,
+                push_buffer_bytes,
+                graph_input_count,
+            )
+            self._world_size = world_size
+            self._pull_config = ConfigResult(NUM_CTA, MAX_THREADS)
+            self.configure_pull(*self._pull_config)  # type: ignore
+
+        @property
+        def world_size(self) -> int:
+            return self._world_size
+
+        def all_reduce(
+            self,
+            input: torch.Tensor,
+            algo: AllReduceAlgo,
+        ) -> tvm_ffi.Tensor:
+            compile_fn = (
+                _jit_custom_all_reduce_push_module
+                if algo.is_push()
+                else _jit_custom_all_reduce_pull_module
+            )
+            module = compile_fn(input.dtype, self._world_size)
+            return module.all_reduce(self, input, algo.shot)
+
+        def config_pull(
+            self, num_blocks: int = -1, num_threads: int = -1
+        ) -> ConfigResult:
+            old_config = self._pull_config
+            num_blocks = num_blocks if num_blocks != -1 else old_config.num_blocks
+            num_threads = num_threads if num_threads != -1 else old_config.num_threads
+            new_config = ConfigResult(num_blocks, num_threads)
+            if new_config != old_config:
+                result = ConfigResult(*self.configure_pull(*new_config))  # type: ignore
+                assert result == self._pull_config
+                self._pull_config = new_config
+            return old_config
+
+        def free(self, tp_cpu_group: torch.distributed.ProcessGroup) -> None:
+            self.free_ipc_handles()  # type: ignore
+            torch.distributed.barrier(group=tp_cpu_group)
+            self.free_storage()  # type: ignore
+
+    return cast(type["CustomAllReduceObj"], CustomAllReduceObjReal)
diff --git a/python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py b/python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py
new file mode 100644
index 000000000..82d5166df
--- /dev/null
+++ b/python/sglang/jit_kernel/benchmark/bench_custom_all_reduce.py
@@ -0,0 +1,377 @@
+"""
+Benchmark JIT custom all-reduce (v2) vs NCCL vs AOT custom all-reduce (v1).
+
+Usage (torchrun required for multi-GPU):
+    torchrun --nproc_per_node=2 bench_custom_all_reduce.py
+    torchrun --nproc_per_node=4 bench_custom_all_reduce.py --dtype float16
+    torchrun --nproc_per_node=8 bench_custom_all_reduce.py --warmup 10 --iters 100
+
+The script initializes all three backends, then benchmarks each over a sweep
+of message sizes.  Results are printed as a comparison table on rank 0.
+"""
+
+import argparse
+import contextlib
+import gc
+import logging
+import os
+from math import isnan
+from typing import Dict, List, Optional
+
+import torch
+import torch.distributed as dist
+
+from sglang.jit_kernel.benchmark.utils import is_in_ci
+
+DTYPE_MAP = {
+    "float16": torch.float16,
+    "bfloat16": torch.bfloat16,
+    "float32": torch.float32,
+}
+
+MESSAGE_SIZES_BYTES = [
+    4 * 1024,  # 4K
+    16 * 1024,  # 16K
+    64 * 1024,  # 64K
+    128 * 1024,  # 128K
+    3 * 64 * 1024,  # 192K
+    4 * 64 * 1024,  # 256K
+    3 * 128 * 1024,  # 384K
+    4 * 128 * 1024,  # 512K
+    5 * 128 * 1024,  # 640K
+    6 * 128 * 1024,  # 768K
+    7 * 128 * 1024,  # 896K
+    1 * 1024 * 1024,  # 1M
+    2 * 1024 * 1024,  # 2M
+    3 * 1024 * 1024,  # 2M
+    4 * 1024 * 1024,  # 4M
+    8 * 1024 * 1024,  # 8M
+    16 * 1024 * 1024,  # 16M
+    32 * 1024 * 1024,  # 32M
+]
+
+
+# ---------------------------------------------------------------------------
+# Backend wrappers - each exposes a uniform interface:
+#   .name          - display name
+#   .capture()     - context manager for CUDA-graph recording
+#   .all_reduce()  - perform an all-reduce and return the result tensor
+# ---------------------------------------------------------------------------
+
+
+class NCCLAllReduceBackend:
+    name = "NCCL"
+
+    def __init__(self, group: dist.ProcessGroup):
+        self.group = group
+
+    def capture(self, register_input: bool):
+        return contextlib.nullcontext()
+
+    def all_reduce(self, tensor: torch.Tensor) -> torch.Tensor:
+        dist.all_reduce(tensor, group=self.group)
+        return tensor
+
+
+class AOTAllReduceBackend:
+    name = "AOT"
+
+    def __init__(self, group: dist.ProcessGroup, device: torch.device):
+        from sglang.srt.distributed.device_communicators.custom_all_reduce import (
+            CustomAllreduce,
+        )
+
+        max_size = max(MESSAGE_SIZES_BYTES)
+        self.comm = CustomAllreduce(group, device, max_size=max_size)
+        if self.comm.disabled:
+            raise RuntimeError("AOT CustomAllreduce is disabled on this system")
+
+    def capture(self, register_input: bool):
+        return self.comm.capture()  # ignore register_input since v1 always requires it
+
+    def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
+        assert self.comm.should_custom_ar(tensor), str(tensor.shape)
+        return self.comm.custom_all_reduce(tensor)
+
+
+class JITAllReduceBackend:
+    name = "JIT"
+
+    def __init__(self, group: dist.ProcessGroup, device: torch.device):
+        from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+            CustomAllReduceV2,
+        )
+
+        max_size = max(MESSAGE_SIZES_BYTES)
+        self.comm = CustomAllReduceV2(group, device, max_pull_size=max_size)
+        if self.comm.disabled:
+            raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
+
+    def capture(self, register_input: bool):
+        return self.comm.capture() if register_input else contextlib.nullcontext()
+
+    def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
+        assert self.comm.should_custom_ar(tensor), str(tensor.shape)
+        return self.comm.custom_all_reduce(tensor)
+
+
+class FlashInferAllReduceBackend:
+    name = "FI"
+
+    def __init__(self, group: dist.ProcessGroup, dtype: torch.dtype):
+        import flashinfer.comm as comm
+
+        rank = torch.distributed.get_rank(group=group)
+        world_size = torch.distributed.get_world_size(group=group)
+        max_size = max(MESSAGE_SIZES_BYTES)
+        hidden_dim = min(MESSAGE_SIZES_BYTES) // 2
+        num_tokens = max_size // hidden_dim
+        self.comm = comm
+        self.hidden_dim = hidden_dim
+        self.workspace = comm.create_allreduce_fusion_workspace(
+            backend="trtllm",
+            world_size=world_size,
+            rank=rank,
+            max_token_num=num_tokens,
+            hidden_dim=hidden_dim,
+            dtype=dtype,
+        )
+
+    def capture(self, *_):
+        return contextlib.nullcontext()
+
+    def all_reduce(self, tensor: torch.Tensor) -> Optional[torch.Tensor]:
+        return self.comm.allreduce_fusion(
+            input=tensor.view(-1, self.hidden_dim),
+            workspace=self.workspace,
+            pattern=self.comm.AllReduceFusionPattern.kAllReduce,
+            launch_with_pdl=True,
+            fp32_acc=True,
+        )
+
+
+# ---------------------------------------------------------------------------
+# Benchmarking helpers
+# ---------------------------------------------------------------------------
+
+
+def parse_args():
+    p = argparse.ArgumentParser(description=__doc__)
+    p.add_argument("--dtype", choices=DTYPE_MAP.keys(), default="bfloat16")
+    p.add_argument("--warmup", type=int, default=5)
+    p.add_argument("--iters", type=int, default=50)
+    p.add_argument("--no-inplace", dest="register_input", action="store_false")
+    return p.parse_args()
+
+
+@torch.inference_mode()
+def bench_one(
+    backend,
+    inp: torch.Tensor,
+    warmup: int,
+    iters: int,
+    group: dist.ProcessGroup,
+    register_input: bool,
+) -> float:
+    """
+    Run *warmup* iterations of all-reduce first.
+    Return the average time for *iters* iterations of all-reduce.
+    """
+    dist.barrier(group=group)
+    for _ in range(warmup):
+        backend.all_reduce(inp)
+    torch.cuda.synchronize()
+
+    # Capture a CUDA graph with *iters* all-reduce calls.
+    inp_batch = torch.stack([inp] * 4)
+    graph = torch.cuda.CUDAGraph()
+    with backend.capture(register_input):
+        with torch.cuda.graph(graph):
+            for i in range(iters):
+                backend.all_reduce(inp_batch[i % 4])
+
+    torch.cuda.synchronize()
+    # Warm up the graph once.
+    graph.replay()
+
+    # Timed replay.
+    start = torch.cuda.Event(enable_timing=True)
+    end = torch.cuda.Event(enable_timing=True)
+    torch.cuda.synchronize()
+    dist.barrier(group=group)
+    graph.replay()  # make the stream busy
+    start.record()
+    graph.replay()
+    end.record()
+    torch.cuda.synchronize()
+    return start.elapsed_time(end) / iters
+
+
+def bench_sweep(
+    backend,
+    sizes_bytes: List[int],
+    dtype: torch.dtype,
+    device: torch.device,
+    warmup: int,
+    iters: int,
+    group: dist.ProcessGroup,
+    register_input: bool,
+) -> Dict[int, float]:
+    """Benchmark one backend over all message sizes."""
+    elem_size = torch.tensor([], dtype=dtype).element_size()
+    results: Dict[int, float] = {}
+    for sz in sizes_bytes:
+        numel = sz // elem_size
+        inp = torch.zeros(numel, dtype=dtype, device=device)
+        try:
+            elapsed_ms = bench_one(backend, inp, warmup, iters, group, register_input)
+            results[sz] = elapsed_ms * 1000  # convert to us per iter
+        except AssertionError:
+            results[sz] = float("nan")
+    return results
+
+
+# ---------------------------------------------------------------------------
+# Result printing
+# ---------------------------------------------------------------------------
+
+
+def print_results(
+    backends: list,
+    all_results: Dict[str, Dict[int, float]],
+    sizes_bytes: List[int],
+) -> None:
+    """Print a comparison table on rank 0."""
+
+    def human_bytes(n: int) -> str:
+        for suffix, unit in [("M", 1 << 20), ("K", 1 << 10)]:
+            if n >= unit and n % unit == 0:
+                return f"{n // unit}{suffix}"
+        return f"{n}B"
+
+    def fmt_us(v: float) -> str:
+        return f"{v:13.1f}" if not isnan(v) else "     n/a"
+
+    names = [b.name for b in backends]
+    nccl_name = "NCCL"
+
+    # Header
+    header_cols = [f"{n:>13}" for n in names]
+    speedup_cols = [f"{n:>13}/NCCL" for n in names if n != nccl_name]
+    header = f"{'Size':>8}  " + "  ".join(header_cols)
+    for sc in speedup_cols:
+        header += f"  {sc}"
+    header += "  "
+    print()
+    print(header)
+    print("-" * len(header))
+
+    # Rows
+    for sz in sizes_bytes:
+        row = f"{human_bytes(sz):>8}"
+        nccl_lat = all_results[nccl_name][sz]
+        for n in names:
+            row += f"  {fmt_us(all_results[n][sz])}"
+        for n in names:
+            if n == nccl_name:
+                continue
+            lat = all_results[n][sz]
+            if not isnan(lat):
+                row += f"  {nccl_lat / lat:17.2f}x"
+            else:
+                row += f"  {'n/a':>17}"
+        print(row)
+
+
+# ---------------------------------------------------------------------------
+# Distributed setup
+# ---------------------------------------------------------------------------
+
+
+def init_distributed():
+    """Initialize distributed groups using torchrun env vars.
+
+    Returns (rank, world_size, device, cpu_group, nccl_group).
+    """
+    import sglang.srt.distributed.parallel_state as ps
+
+    local_rank = int(os.environ.get("LOCAL_RANK", "0"))
+    world_size = int(os.environ.get("WORLD_SIZE", "1"))
+    rank = local_rank
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+    torch.cuda.set_stream(torch.cuda.Stream())  # use a non-default stream
+
+    torch.distributed.init_process_group(backend="gloo")
+    ps._WORLD = coord = ps.init_world_group(
+        ranks=list(range(world_size)),
+        local_rank=local_rank,
+        backend="nccl",
+    )
+
+    cpu_group = coord.cpu_group
+    nccl_group = coord.device_group
+    assert nccl_group is not None
+    return rank, world_size, device, cpu_group, nccl_group
+
+
+# ---------------------------------------------------------------------------
+# Main
+# ---------------------------------------------------------------------------
+
+
+def main():
+    logging.basicConfig(level=logging.WARNING)
+    args = parse_args()
+    dtype = DTYPE_MAP[args.dtype]
+
+    rank, world_size, device, cpu_group, nccl_group = init_distributed()
+
+    # Instantiate backends.
+    backends = [
+        NCCLAllReduceBackend(nccl_group),
+        JITAllReduceBackend(cpu_group, device),
+    ]
+    if world_size in [2, 4, 6, 8]:
+        backends.insert(1, AOTAllReduceBackend(cpu_group, device))
+    if world_size in [2, 4, 8]:
+        backends.append(FlashInferAllReduceBackend(cpu_group, dtype))
+
+    # Run benchmarks.
+    all_results: Dict[str, Dict[int, float]] = {}
+    torch.cuda.synchronize()
+    for backend in backends:
+        if rank == 0:
+            print(f"Benchmarking {backend.name} ...")
+        all_results[backend.name] = bench_sweep(
+            backend,
+            MESSAGE_SIZES_BYTES,
+            dtype,
+            device,
+            args.warmup,
+            args.iters,
+            cpu_group,
+            args.register_input,
+        )
+
+    # Aggregate across ranks (use max to reflect the slowest rank).
+    for name in list(all_results):
+        for sz in MESSAGE_SIZES_BYTES:
+            val = all_results[name].get(sz)
+            if val is None:
+                continue
+            t = torch.tensor([val], dtype=torch.float64, device=device)
+            dist.all_reduce(t, op=dist.ReduceOp.MAX, group=nccl_group)
+            all_results[name][sz] = t.item()
+
+    # Print results on rank 0.
+    if rank == 0:
+        print_results(backends, all_results, MESSAGE_SIZES_BYTES)
+
+    del backends, all_results
+    gc.collect()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__" and not is_in_ci():
+    main()
diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh
new file mode 100644
index 000000000..dc5f5beea
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_base.cuh
@@ -0,0 +1,27 @@
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/custom_all_reduce.cuh>
+
+#include <cstdint>
+#include <cstring>
+
+inline void register_custom_all_reduce() {
+  namespace refl = tvm::ffi::reflection;
+  using Class = host::distributed::CustomAllReduceBase;
+  refl::ObjectDef<Class>()
+      .def(refl::init<uint32_t, uint32_t, uint32_t, uint32_t, int64_t, int64_t, int64_t>(), "__init__")
+      .def("share_storage", &Class::share_storage)
+      .def("share_graph_inputs", &Class::share_graph_inputs)
+      .def("post_init", &Class::post_init)
+      .def("register_inputs", &Class::register_inputs)
+      .def("set_cuda_graph_capture", &Class::set_cuda_graph_capture)
+      .def("free_ipc_handles", &Class::free_ipc_handles)
+      .def("free_storage", &Class::free_storage)
+      .def("configure_pull", &Class::configure_pull);
+}
diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh
new file mode 100644
index 000000000..e8837af4c
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_pull.cuh
@@ -0,0 +1,205 @@
+// Partially migrated from AOT kernel:
+// https://github.com/sgl-project/sglang/blob/v0.5.9/sgl-kernel/csrc/allreduce/custom_all_reduce.cu
+// Which was originally adapted from:
+// https://github.com/vllm-project/vllm/blob/v0.8.2/csrc/custom_all_reduce.cu
+// We redesign the controller interface to minimize control plane traffic,
+// and fuse the reduce-scatter and broadcast in the 2-shot all reduce
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/common.cuh>
+#include <sgl_kernel/distributed/custom_all_reduce.cuh>
+
+#include <bit>
+#include <cstdint>
+#include <cstring>
+
+namespace {
+
+using device::distributed::PullController;
+using host::distributed::AllReduceData;
+using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
+
+struct AllReduceParams {
+  void* __restrict__ output;
+  uint32_t rank;
+  uint32_t num_items;  // NOTE: support at most 4G, but that's too much
+};
+
+[[maybe_unused]]
+SGL_DEVICE void prefetch_uniform_ptr(const void* ptr) {
+  asm volatile("prefetchu.L1 [%0];" ::"l"(ptr) : "memory");
+}
+
+#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
+
+template <bool kBroadcast, typename DType, uint32_t kNumGPU>
+SGL_DEVICE void all_reduce_impl(const AllReduceParams& params, DType* (&input)[kNumGPU]) {
+  using namespace device;
+
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  using DType2 = packed_t<DType>;
+  using Storage = AlignedVector<DType2, kVecSize>;
+  const auto& [output, rank, num_items] = params;
+
+  for (auto i = blockIdx.x;; i += gridDim.x) {
+    const auto offset = i * blockDim.x + threadIdx.x;
+    if (offset * kVecSize * 2 >= num_items) break;
+    Storage storage[kNumGPU];
+
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      storage[i].load(input[i], offset);
+    }
+    const Storage result = distributed::reduce_impl(storage);
+    if constexpr (kBroadcast) {
+#pragma unroll
+      for (uint32_t i = 0; i < kNumGPU; ++i) {
+        result.store(input[i], offset);
+      }
+    } else {
+      result.store(output, offset);
+    }
+  }
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+CUSTOM_AR_KERNEL void all_reduce_one_shot_kernel(
+    const AllReduceData* __restrict__ data,
+    const AllReduceParams __grid_constant__ params,
+    const PullController __grid_constant__ ctrl) {
+  /// NOTE: we assume the data array is ready before the previous kernel
+  DType* input[kNumGPU];
+  prefetch_uniform_ptr(data);
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i)
+    input[i] = static_cast<DType*>(data->input[i]);
+  device::PDLWaitPrimary<kUsePDL>();
+
+  ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
+  all_reduce_impl</*kBroadcast=*/false>(params, input);
+
+  device::PDLTriggerSecondary<kUsePDL>();
+  ctrl.sync</*kFence=*/0, /*kStart=*/0>(params.rank, kNumGPU);
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+CUSTOM_AR_KERNEL void all_reduce_two_shot_kernel(
+    const AllReduceData* __restrict__ data,
+    const AllReduceParams __grid_constant__ params,
+    const PullController __grid_constant__ ctrl) {
+  // get the range of this rank
+  using device::kWarpThreads, device::div_ceil;
+
+  prefetch_uniform_ptr(data);
+  DType* input[kNumGPU];
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i)
+    input[i] = static_cast<DType*>(data->input[i]);
+
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  const uint32_t num_items = params.num_items;
+  const uint32_t total_vec = num_items / (kVecSize * 2);  // must be divisible here
+  const uint32_t vec_per_rank = div_ceil(div_ceil(total_vec, kNumGPU), kWarpThreads) * kWarpThreads;
+  const uint32_t local_vec_start = min(params.rank * vec_per_rank, total_vec);
+  const uint32_t local_vec_finish = min(local_vec_start + vec_per_rank, total_vec);
+  const uint32_t local_start = local_vec_start * kVecSize * 2;
+  const uint32_t local_length = (local_vec_finish - local_vec_start) * kVecSize * 2;
+  const auto local_params = AllReduceParams{
+      .output = nullptr,  // this is not used for 2-shot all reduce
+      .rank = params.rank,
+      .num_items = local_length,
+  };
+
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i)
+    input[i] += local_start;
+
+  device::PDLWaitPrimary<kUsePDL>();
+
+  ctrl.sync</*kFence=*/0, /*kStart=*/1>(params.rank, kNumGPU);
+  all_reduce_impl</*kBroadcast=*/true>(local_params, input);
+
+  device::PDLTriggerSecondary<kUsePDL>();
+  ctrl.sync</*kFence=*/1, /*kStart=*/0>(params.rank, kNumGPU);
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+struct CustomAllReducePull : public CustomAllReduceBase {
+  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  static constexpr auto one_shot_kernel = all_reduce_one_shot_kernel<DType, kNumGPU, kUsePDL>;
+  static constexpr auto two_shot_kernel = all_reduce_two_shot_kernel<DType, kNumGPU, kUsePDL>;
+  static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
+
+  tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
+    using namespace host;
+    const bool use_2shot = (shot == 2);
+    const auto device = input.device();
+    const auto input_ptr = input.data_ptr();
+    const auto buffer_ptr = get_pull_buffer(m_storage);
+    const auto num_items_int64 = input.numel();
+    const auto num_items = static_cast<uint32_t>(num_items_int64);
+    const auto items_per_block = m_cta_size * kVecSize * 2;
+    const auto needed_blocks = div_ceil(num_items, items_per_block);
+    const auto num_blocks = std::min(needed_blocks, m_num_cta);
+    const auto kernel = use_2shot ? two_shot_kernel : one_shot_kernel;
+    // only 1-shot + graph capture need extra output buffer
+    const auto output = (m_is_graph_capturing && !use_2shot) ? ffi::empty_like(input) : input;
+    const auto params = AllReduceParams{
+        .output = use_2shot ? nullptr : output.data_ptr(),
+        .rank = m_rank,
+        .num_items = num_items,
+    };
+
+    RuntimeCheck(input.IsContiguous(), "Input tensor must be contiguous");
+    RuntimeCheck(m_num_gpu == kNumGPU, "Mismatch GPU count");
+    RuntimeCheck(shot == 1 || shot == 2, "Invalid shot count: ", shot);
+    RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
+    RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
+    RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
+    RuntimeCheck(m_pull_ctrl.has_value(), "Controller is not initialized");
+    RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
+
+    const auto& ctrl = *m_pull_ctrl;
+    const auto stream = LaunchKernel::resolve_device(device);
+    auto launch = LaunchKernel{num_blocks, m_cta_size, stream};
+    launch.enable_pdl(kUsePDL);
+    const auto check_capturing = [&] {
+      if (!m_is_graph_capturing) return false;  // override to avoid cudaRT call overhead
+      cudaStreamCaptureStatus status;
+      RuntimeDeviceCheck(cudaStreamIsCapturing(stream, &status));
+      return status == cudaStreamCaptureStatusActive;
+    };
+    if (check_capturing()) {
+      // no-op if not really capturing, we're in a dummy run
+      const auto data_ptr = allocate_graph_capture_input(input_ptr);
+      /// NOTE: we assume when the graph is replayed, the data_ptr should be ready
+      launch(kernel, data_ptr, params, ctrl);
+    } else {
+      // 1.copy the input to the buffer
+      const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items);
+      RuntimeCheck(input_bytes <= m_pull_buffer_bytes, "Input is too large, num items: ", num_items);
+      RuntimeDeviceCheck(cudaMemcpyAsync(buffer_ptr, input_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
+      // 2. launch the all reduce kernel
+      const auto data_ptr = get_data_ptr();  // use default buffer
+      launch(kernel, data_ptr, params, ctrl);
+      if (use_2shot) {  // 3. copy the reduced result back to the output, because 2-shot doesn't write to output
+        RuntimeDeviceCheck(cudaMemcpyAsync(input_ptr, buffer_ptr, input_bytes, cudaMemcpyDeviceToDevice, stream));
+      }
+    }
+    return output;
+  }
+};
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
+  using Impl = CustomAllReducePull<DType, kNumGPU, kUsePDL>;
+  return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
+}
+
+}  // namespace
diff --git a/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh
new file mode 100644
index 000000000..c4523c27e
--- /dev/null
+++ b/python/sglang/jit_kernel/csrc/distributed/custom_all_reduce_push.cuh
@@ -0,0 +1,253 @@
+// Partially adapted from:
+// https://github.com/flashinfer-ai/flashinfer/blob/v0.6.4/include/flashinfer/comm/trtllm_allreduce_fusion.cuh
+// We simplify the lamport design and minimize the ring buffer count (from 3 -> 2)
+#include <sgl_kernel/ffi.h>
+#include <sgl_kernel/tensor.h>
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/type.cuh>
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/common.cuh>
+#include <sgl_kernel/distributed/custom_all_reduce.cuh>
+
+#include <cstdint>
+#include <cstring>
+
+namespace {
+
+using device::distributed::PushController;
+using host::distributed::CustomAllReduceBase, host::distributed::CustomAllReduceRef;
+
+struct AllReducePushData {
+  void* __restrict__ buffer[device::distributed::kMaxNumGPU];
+  const void* input;
+  void* output;
+  uint32_t rank;
+  uint32_t num_items;
+  uint32_t buffer_bytes;
+  uint32_t epoch_bytes;
+};
+
+#define CUSTOM_AR_KERNEL __global__ __launch_bounds__(1024, 1)
+
+template <typename T>
+struct fp_trait {};
+
+// TODO: support more dtypes
+template <>
+struct fp_trait<bf16_t> {
+  using type = uint16_t;
+  [[maybe_unused]]
+  static constexpr uint16_t pos_zero = 0x0000u;
+  [[maybe_unused]]
+  static constexpr uint16_t neg_zero = 0x8000u;
+};
+
+template <>
+struct fp_trait<fp16_t> {
+  using type = uint16_t;
+  [[maybe_unused]]
+  static constexpr uint16_t pos_zero = 0x0000u;
+  [[maybe_unused]]
+  static constexpr uint16_t neg_zero = 0x8000u;
+};
+
+template <>
+struct fp_trait<float> {
+  using type = uint32_t;
+  [[maybe_unused]]
+  static constexpr uint32_t pos_zero = 0x00000000u;
+  [[maybe_unused]]
+  static constexpr uint32_t neg_zero = 0x80000000u;
+};
+
+template <typename DType>
+SGL_DEVICE void clear_pos_zero(DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<typename Trait::type*>(&val);
+  if (*ptr == Trait::pos_zero) *ptr = Trait::neg_zero;
+}
+
+template <typename DType>
+SGL_DEVICE bool is_pos_zero(const DType& val) {
+  using Trait = fp_trait<DType>;
+  const auto ptr = reinterpret_cast<const typename Trait::type*>(&val);
+  return *ptr == Trait::pos_zero;
+}
+
+template <typename DType>
+SGL_DEVICE DType get_pos_zero() {
+  using Trait = fp_trait<DType>;
+  const auto value = Trait::pos_zero;
+  return *reinterpret_cast<const DType*>(&value);
+}
+
+template <typename T>
+SGL_DEVICE void ld_global_volatile_16B(T& x, const void* addr, int64_t offset) {
+  static_assert(alignof(T) == 16 && sizeof(T) == 16);
+  addr = device::pointer::offset<T>(addr, offset);
+  uint4 val;
+  asm volatile("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];"
+               : "=r"(val.x), "=r"(val.y), "=r"(val.z), "=r"(val.w)
+               : "l"(addr));
+  x = *reinterpret_cast<const T*>(&val);
+}
+
+template <typename T>
+SGL_DEVICE void st_global_volatile_16B(const T& x, void* addr, int64_t offset) {
+  static_assert(alignof(T) == 16 && sizeof(T) == 16);
+  const uint4 val = *reinterpret_cast<const uint4*>(&x);
+  addr = device::pointer::offset<T>(addr, offset);
+  asm volatile(
+      "st.volatile.global.v4.b32 [%4], {%0, %1, %2, %3};" ::"r"(val.x), "r"(val.y), "r"(val.z), "r"(val.w), "l"(addr));
+}
+
+template <typename DType, uint32_t kNumGPU>
+SGL_DEVICE void push_impl(DType* (&push_buf)[kNumGPU], const void* data, uint32_t num_items) {
+  using namespace device;
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  using Storage = AlignedVector<packed_t<DType>, kVecSize>;
+
+  for (auto i = blockIdx.x;; i += gridDim.x) {
+    const auto offset = i * blockDim.x + threadIdx.x;
+    if (offset * kVecSize * 2 >= num_items) break;
+    Storage vec;
+    vec.load(data, offset);
+#pragma unroll
+    for (uint32_t j = 0; j < kVecSize; ++j) {
+      clear_pos_zero(vec[j].x);
+      clear_pos_zero(vec[j].y);
+    }
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      st_global_volatile_16B(vec, push_buf[i], offset);
+    }
+  }
+}
+
+template <typename DType, uint32_t kNumGPU>
+SGL_DEVICE void poll_impl(DType* (&poll_buf)[kNumGPU], void* data, uint32_t num_items) {
+  using namespace device;
+  constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  using Storage = AlignedVector<packed_t<DType>, kVecSize>;
+
+  for (auto i = blockIdx.x;; i += gridDim.x) {
+    const auto offset = i * blockDim.x + threadIdx.x;
+    if (offset * kVecSize * 2 >= num_items) break;
+    Storage storage[kNumGPU];
+
+    while (true) {
+      bool has_pos_zero = false;
+#pragma unroll
+      for (uint32_t i = 0; i < kNumGPU; ++i) {
+        ld_global_volatile_16B(storage[i], poll_buf[i], offset);
+#pragma unroll
+        for (auto j = 0; j < kVecSize; ++j) {
+          has_pos_zero |= is_pos_zero(storage[i][j].x);
+          has_pos_zero |= is_pos_zero(storage[i][j].y);
+        }
+      }
+      if (!has_pos_zero) break;
+    }
+
+    const Storage result = distributed::reduce_impl(storage);
+    result.store(data, offset);
+
+    Storage pos_zeros;
+    pos_zeros.fill({get_pos_zero<DType>(), get_pos_zero<DType>()});
+#pragma unroll
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      pos_zeros.store(poll_buf[i], offset);
+    }
+  }
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+CUSTOM_AR_KERNEL void all_reduce_one_shot_push_kernel(
+    const AllReducePushData __grid_constant__ params,  //
+    const PushController __grid_constant__ ctrl) {
+  using namespace device;
+
+  const auto [buffer, input, output, rank, num_items, buffer_bytes, epoch_bytes] = params;
+
+  PDLWaitPrimary<kUsePDL>();
+
+  // Phase 1: Push data from input to all ranks' buffers
+  const auto epoch_offset = ctrl.epoch() * epoch_bytes;
+  DType* push_buf[kNumGPU];
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i) {
+    push_buf[i] = static_cast<DType*>(pointer::offset(buffer[i], rank * buffer_bytes, epoch_offset));
+  }
+  push_impl(push_buf, input, num_items);
+
+  PDLTriggerSecondary<kUsePDL>();
+
+  // Phase 2: Poll local data
+  DType* poll_buf[kNumGPU];
+#pragma unroll
+  for (uint32_t i = 0; i < kNumGPU; ++i) {
+    poll_buf[i] = static_cast<DType*>(pointer::offset(buffer[rank], i * buffer_bytes, epoch_offset));
+  }
+  poll_impl(poll_buf, output, num_items);
+  ctrl.exit();
+}
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+struct CustomAllReducePush : public CustomAllReduceBase {
+  static constexpr uint32_t kVecSize = 16 / (sizeof(DType) * 2);
+  static_assert(kNumGPU <= device::distributed::kMaxNumGPU, "kNumGPU exceeds the maximum supported GPUs");
+
+  tvm::ffi::Tensor all_reduce(tvm::ffi::Tensor input, int shot) {
+    using namespace host;
+    const auto device = input.device();
+    const auto input_ptr = input.data_ptr();
+    const auto num_items_int64 = input.numel();
+    const auto num_items = static_cast<uint32_t>(num_items_int64);
+    const auto num_blocks = m_max_num_cta_push;  // must be constant to ensure correctness
+    const auto num_threads = [&] {
+      for (const auto t : {128u, 256u, 512u}) {
+        if (t * num_blocks * 2 * kVecSize >= num_items) return t;
+      }
+      return 1024u;
+    }();
+    const auto output = input;
+    AllReducePushData params;
+    for (uint32_t i = 0; i < kNumGPU; ++i) {
+      params.buffer[i] = get_push_buffer(m_peer_storage[i]);
+    }
+    params.input = input_ptr;
+    params.output = input_ptr;
+    params.rank = m_rank;
+    params.num_items = num_items;
+    params.buffer_bytes = m_push_buffer_bytes;
+    params.epoch_bytes = kNumGPU * params.buffer_bytes;
+
+    RuntimeCheck(input.IsContiguous(), "Input must be contiguous");
+    RuntimeCheck(m_num_gpu == kNumGPU, "Number of GPUs mismatch");
+    RuntimeCheck(device.device_type == kDLCUDA, "Only CUDA device is supported");
+    RuntimeCheck(is_type<DType>(input.dtype()), "Input dtype mismatch");
+    RuntimeCheck(std::bit_cast<intptr_t>(input_ptr) % 16 == 0, "Input pointer is not properly aligned");
+    RuntimeCheck(m_push_ctrl.has_value(), "Controller is not initialized");
+    RuntimeCheck(shot == 1, "Push all-reduce only supports 1-shot, got: ", shot);
+    RuntimeCheck(static_cast<int64_t>(num_items) == num_items_int64, "Number of items exceeds 4G limit");
+
+    const auto input_bytes = static_cast<int64_t>(sizeof(DType) * num_items_int64);
+    RuntimeCheck(input_bytes <= m_push_buffer_bytes, "Input is too large, num items: ", num_items);
+
+    const auto kernel = all_reduce_one_shot_push_kernel<DType, kNumGPU, kUsePDL>;
+    LaunchKernel(num_blocks, num_threads, device)  //
+        .enable_pdl(kUsePDL)(kernel, params, *m_push_ctrl);
+    return output;
+  }
+};
+
+template <typename DType, uint32_t kNumGPU, bool kUsePDL>
+tvm::ffi::Tensor custom_all_reduce(CustomAllReduceRef obj, tvm::ffi::Tensor input, int shot) {
+  using Impl = CustomAllReducePush<DType, kNumGPU, kUsePDL>;
+  return static_cast<Impl&>(*obj.get()).all_reduce(input, shot);
+}
+
+}  // namespace
diff --git a/python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh b/python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh
index 15b52d81c..1a88ad02b 100644
--- a/python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh
+++ b/python/sglang/jit_kernel/csrc/gemm/marlin/marlin.cuh
@@ -40,8 +40,6 @@ struct Vec {
 
 using I4 = Vec<int, 4>;
 
-using host::div_ceil;
-
 #if defined(__CUDA_ARCH__) && __CUDA_ARCH__ < 800
 // No support for async
 #else
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh
new file mode 100644
index 000000000..c9711c5b6
--- /dev/null
+++ b/python/sglang/jit_kernel/include/sgl_kernel/distributed/common.cuh
@@ -0,0 +1,114 @@
+#pragma once
+#include <sgl_kernel/utils.cuh>
+
+namespace device::distributed {
+
+inline constexpr uint32_t kMaxNumGPU = 8;
+
+struct alignas(128) Semaphore {
+ public:
+  constexpr Semaphore() : m_flag(0), m_counter(0) {}
+
+  template <bool kFence>
+  SGL_DEVICE uint32_t get() const {
+    uint32_t val;
+    if constexpr (kFence) {
+      asm volatile("ld.acquire.sys.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag));
+    } else {
+      asm volatile("ld.volatile.global.u32 %0, [%1];" : "=r"(val) : "l"(&m_flag));
+    }
+    return val;
+  }
+
+  template <bool kFence>
+  SGL_DEVICE uint32_t add(uint32_t val) {
+    uint32_t old_val;
+    if constexpr (kFence) {
+      asm volatile("atom.release.sys.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val));
+    } else {
+      asm volatile("atom.global.add.u32 %0, [%1], %2;" : "=r"(old_val) : "l"(&m_flag), "r"(val));
+    }
+    return old_val;
+  }
+
+  // Only called by the owning GPU - plain load is sufficient
+  SGL_DEVICE uint32_t get_counter() const {
+    return m_counter;
+  }
+
+  // Only called by the owning GPU - plain store is sufficient
+  SGL_DEVICE void set_counter(uint32_t val) {
+    m_counter = val;
+  }
+
+ private:
+  uint32_t m_flag;
+  uint32_t m_counter;
+};
+
+struct PullController {
+ public:
+  PullController(void** signals, uint32_t num_gpu) {
+    for (uint32_t i = 0; i < num_gpu; ++i) {
+      m_signals[i] = static_cast<Semaphore*>(signals[i]);
+    }
+  }
+
+  /// Synchronize all GPUs.
+  /// When kFence is true, establishes happens-before across GPUs using
+  /// release/acquire semantics, ensuring prior writes are visible system-wide.
+  template <bool kFence, bool kStart>
+  SGL_DEVICE void sync(uint32_t rank, uint32_t num_gpu) const {
+    // For fenced sync: ensure all threads in this block have completed their writes,
+    // so the signaling thread's release carries them transitively.
+    static_assert(!(kFence && kStart), "Start stage does not need to wait fence");
+    if constexpr (kFence || !kStart) __syncthreads();
+    constexpr auto kStage = kStart ? 1 : 2;
+    const auto warp_id = threadIdx.x / kWarpThreads;
+    const auto lane_id = threadIdx.x % kWarpThreads;
+    if (lane_id == 0 && warp_id < num_gpu) {
+      auto& signal = m_signals[warp_id][blockIdx.x];
+      signal.add<kFence>(1);
+      if (warp_id == rank) {
+        const auto target = num_gpu * kStage;
+        /// NOTE: correctness here:
+        /// - base is only read/updated locally by the owning GPU
+        const auto base = signal.get_counter();
+        while (signal.get<kFence>() - base < target)
+          ;
+        if constexpr (!kStart) {
+          signal.set_counter(base + target);
+        }
+      }
+    }
+    if constexpr (kStart) __syncthreads();
+  }
+
+ private:
+  Semaphore* __restrict__ m_signals[kMaxNumGPU];
+};
+
+struct PushController {
+ public:
+  static constexpr int64_t kNumStages = 2;
+
+  PushController(void* ptr) : m_local_signal(static_cast<Semaphore*>(ptr)) {}
+
+  SGL_DEVICE uint32_t epoch() const {
+    return m_local_signal[blockIdx.x].get_counter();
+  }
+
+  SGL_DEVICE void exit() const {
+    __syncthreads();
+    if (threadIdx.x == 0) {
+      auto& signal = m_local_signal[blockIdx.x];
+      const auto epoch = signal.get_counter();
+      signal.set_counter((epoch + 1) % kNumStages);
+    }
+  }
+
+ private:
+  Semaphore* m_local_signal;
+};
+
+}  // namespace device::distributed
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh
new file mode 100644
index 000000000..aeaa58b84
--- /dev/null
+++ b/python/sglang/jit_kernel/include/sgl_kernel/distributed/custom_all_reduce.cuh
@@ -0,0 +1,349 @@
+#pragma once
+#include <sgl_kernel/utils.h>
+
+#include <sgl_kernel/utils.cuh>
+#include <sgl_kernel/vec.cuh>
+
+#include <sgl_kernel/distributed/common.cuh>
+
+#include <tvm/ffi/container/array.h>
+#include <tvm/ffi/container/tuple.h>
+#include <tvm/ffi/reflection/registry.h>
+
+#include <array>
+#include <cstdint>
+#include <cstring>
+#include <functional>
+#include <numeric>
+#include <optional>
+#include <span>
+#include <unordered_map>
+#include <vector>
+
+namespace host::distributed {
+
+using device::distributed::PullController, device::distributed::PushController;
+
+struct AllReduceData {
+  constexpr AllReduceData() {}
+  void* __restrict__ input[device::distributed::kMaxNumGPU];
+};
+
+using ExternHandle = tvm::ffi::Array<char>;
+
+inline ExternHandle to_extern_handle(void* ptr) {
+  ExternHandle array;
+  cudaIpcMemHandle_t handle;
+  RuntimeDeviceCheck(cudaIpcGetMemHandle(&handle, ptr));
+  for (size_t i = 0; i < sizeof(handle); ++i) {
+    array.push_back(handle.reserved[i]);
+  }
+  return array;
+}
+
+inline void* from_extern_handle(const ExternHandle& array) {
+  cudaIpcMemHandle_t handle;
+  RuntimeCheck(array.size() == sizeof(handle), "Invalid IPC handle size: ", array.size());
+  for (size_t i = 0; i < sizeof(handle); ++i) {
+    handle.reserved[i] = array[i];
+  }
+  void* ptr;
+  RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess));
+  return ptr;
+}
+
+struct HandleHash {
+  std::size_t operator()(const cudaIpcMemHandle_t& handle) const {
+    return std::hash<std::string_view>{}({handle.reserved, sizeof(handle.reserved)});
+  }
+};
+
+struct HandleEqual {
+  bool operator()(const cudaIpcMemHandle_t& a, const cudaIpcMemHandle_t& b) const {
+    return std::memcmp(a.reserved, b.reserved, sizeof(a.reserved)) == 0;
+  }
+};
+
+/**
+ * \brief The control plane of the custom all-reduce implementation.
+ * It manages the internal state and synchronization of the participating GPUs.
+ */
+struct CustomAllReduceBase : public tvm::ffi::Object {
+ public:
+  TVM_FFI_DECLARE_OBJECT_INFO_FINAL("sgl.CustomAllReduce", CustomAllReduceBase, tvm::ffi::Object);
+
+  static constexpr bool _type_mutable = true;
+  using InputPair = tvm::ffi::Tuple<int64_t, ExternHandle>;  // (offset, ipc handle)
+
+  CustomAllReduceBase(
+      uint32_t rank,
+      uint32_t num_gpu,
+      uint32_t max_num_cta_pull,
+      uint32_t max_num_cta_push,
+      int64_t pull_buffer_size,
+      int64_t push_buffer_size,
+      int64_t graph_buffer_count)
+      : m_pull_buffer_bytes(pull_buffer_size),
+        m_push_buffer_bytes(push_buffer_size),
+        m_graph_buffer_count(graph_buffer_count),
+        m_rank(rank),
+        m_num_gpu(num_gpu),
+        m_max_num_cta_pull(max_num_cta_pull),
+        m_max_num_cta_push(max_num_cta_push),
+        // default config for pull kernel, can be updated by `configure()`
+        m_num_cta(max_num_cta_pull),
+        m_cta_size(256) {
+    RuntimeDeviceCheck(cudaMalloc(&m_storage, storage_bytes()));
+    RuntimeCheck(rank < num_gpu, "Invalid rank: ", rank);
+    const int64_t kU32Max = static_cast<int64_t>(std::numeric_limits<uint32_t>::max());
+    const int64_t push_buffer_size_all = push_all_ranks_bytes();
+    RuntimeCheck(pull_buffer_size <= kU32Max, "Buffer size is too large: ", pull_buffer_size);
+    RuntimeCheck(push_buffer_size_all <= kU32Max, "Push buffer size is too large: ", push_buffer_size_all);
+  }
+
+  ExternHandle share_storage() {
+    return to_extern_handle(m_storage);
+  }
+
+  tvm::ffi::Array<InputPair> share_graph_inputs() {
+    tvm::ffi::Array<InputPair> result;
+    const auto new_inputs_count = registered_count() - m_cum_registered_count;
+    RuntimeCheck(new_inputs_count >= 0, "Invalid new count: ", new_inputs_count);
+    result.reserve(new_inputs_count);
+    std::unordered_map<void*, ExternHandle> ipc_cache;
+    const auto get_handle = [&](void* ptr) -> ExternHandle {
+      const auto it = ipc_cache.find(ptr);
+      if (it != ipc_cache.end()) return it->second;
+      const auto handle = to_extern_handle(ptr);
+      ipc_cache.try_emplace(ptr, handle);
+      return handle;
+    };
+    for (const auto ptr : std::span(m_graph_capture_inputs).subspan(m_cum_registered_count)) {
+      // note: must share the base address of each allocation, or we get wrong address
+      void* base_ptr;
+      const auto cu_result = cuPointerGetAttribute(&base_ptr, CU_POINTER_ATTRIBUTE_RANGE_START_ADDR, (CUdeviceptr)ptr);
+      RuntimeCheck(cu_result == CUDA_SUCCESS, "failed to get pointer attr");
+      const auto offset = reinterpret_cast<char*>(ptr) - reinterpret_cast<char*>(base_ptr);
+      result.push_back(InputPair{offset, get_handle(base_ptr)});
+    }
+    return result;
+  }
+
+  void post_init(tvm::ffi::Array<ExternHandle> ipc_storages) {
+    RuntimeCheck(ipc_storages.size() == m_num_gpu, "Invalid array size: ", ipc_storages.size());
+    m_peer_storage.resize(m_num_gpu);
+    for (const auto i : irange(m_num_gpu)) {
+      if (i == m_rank) {
+        m_peer_storage[i] = m_storage;
+      } else {
+        m_peer_storage[i] = from_extern_handle(ipc_storages[i]);
+      }
+    }
+
+    // set signal buffer to zero
+    const auto pull_signal = get_pull_signal(m_storage);
+    RuntimeDeviceCheck(cudaMemset(pull_signal, 0, pull_signal_bytes()));
+
+    // update the pull controller and data pointer
+    RuntimeCheck(!m_pull_ctrl.has_value(), "Controller is already initialized");
+    m_pull_ctrl.emplace(m_peer_storage.data(), m_num_gpu);
+    AllReduceData data;
+    for (const auto i : irange(m_num_gpu)) {
+      data.input[i] = get_pull_buffer(m_peer_storage[i]);
+    }
+    const auto default_data_ptr = get_data_ptr();
+    RuntimeDeviceCheck(cudaMemcpy(default_data_ptr, &data, sizeof(AllReduceData), cudaMemcpyHostToDevice));
+
+    // update the push controller and data pointer
+    RuntimeCheck(!m_push_ctrl.has_value(), "Controller is already initialized");
+    const auto push_signal = get_push_signal(m_storage);
+    RuntimeDeviceCheck(cudaMemset(push_signal, 0, push_signal_bytes()));
+    m_push_ctrl.emplace(push_signal);
+    const auto push_buffer = get_push_buffer(m_storage);
+    RuntimeDeviceCheck(cudaMemset(push_buffer, 0, push_all_ranks_bytes()));
+  }
+
+  void register_inputs(tvm::ffi::Array<tvm::ffi::Array<InputPair>> ipc_graph_inputs) {
+    RuntimeCheck(ipc_graph_inputs.size() == m_num_gpu);
+    const auto new_registered_count = registered_count() - m_cum_registered_count;
+    RuntimeCheck(new_registered_count >= 0, "Invalid registered count: ", new_registered_count);
+    if (new_registered_count == 0) return;  // avoid `m_get_data_ptr()` out-of-bounds
+    std::vector<AllReduceData> data;
+    data.resize(new_registered_count);
+    const auto open_cached = [&](const ExternHandle& h) -> void* {
+      RuntimeCheck(h.size() == sizeof(cudaIpcMemHandle_t), "Invalid IPC handle size: ", h.size());
+      cudaIpcMemHandle_t handle;
+      for (size_t i = 0; i < sizeof(handle); ++i)
+        handle.reserved[i] = h[i];
+      const auto [it, success] = m_ipc_cache.try_emplace(handle, nullptr);
+      if (success) {
+        void* ptr;
+        RuntimeDeviceCheck(cudaIpcOpenMemHandle(&ptr, handle, cudaIpcMemLazyEnablePeerAccess));
+        it->second = ptr;
+      }
+      return it->second;
+    };
+    for (const auto i : irange(ipc_graph_inputs.size())) {
+      const auto& array = ipc_graph_inputs[i];
+      RuntimeCheck(int64_t(array.size()) == new_registered_count);
+      if (i == m_rank) {
+        for (const auto j : irange(new_registered_count)) {
+          data[j].input[i] = m_graph_capture_inputs[m_cum_registered_count + j];
+        }
+      } else {
+        for (const auto j : irange(new_registered_count)) {
+          /// NOTE: structural binding will cause intern compiler error...
+          const auto elem = array[j];
+          const auto offset = get<0>(elem);
+          const auto ipc_handle = get<1>(elem);
+          data[j].input[i] = pointer::offset(open_cached(ipc_handle), offset);
+        }
+      }
+    }
+
+    const auto new_registered_bytes = sizeof(AllReduceData) * new_registered_count;
+    const auto dst_ptr = get_data_ptr(m_cum_registered_count);
+    m_cum_registered_count += new_registered_count;
+    RuntimeDeviceCheck(cudaMemcpy(dst_ptr, data.data(), new_registered_bytes, cudaMemcpyHostToDevice));
+  }
+
+  void set_cuda_graph_capture(bool enabled) {
+    m_is_graph_capturing = enabled;
+  }
+
+  void free_ipc_handles() {
+    for (const auto& pair : m_ipc_cache) {
+      host::RuntimeDeviceCheck(cudaIpcCloseMemHandle(pair.second));
+    }
+    m_ipc_cache.clear();
+  }
+
+  void free_storage() {
+    host::RuntimeDeviceCheck(cudaFree(m_storage));
+    m_storage = nullptr;
+  }
+
+  tvm::ffi::Tuple<uint32_t, uint32_t> configure_pull(uint32_t num_cta, uint32_t cta_size) {
+    using host::RuntimeCheck;
+    const auto min_cta_size = m_num_gpu * device::kWarpThreads;
+    RuntimeCheck(num_cta > 0 && num_cta <= m_max_num_cta_pull, "Invalid number of CTAs: ", num_cta);
+    RuntimeCheck(cta_size >= min_cta_size, "Block size must be at least ", min_cta_size);
+    const auto old_num_cta = m_num_cta;
+    const auto old_block_size = m_cta_size;
+    m_num_cta = num_cta;
+    m_cta_size = cta_size;
+    return tvm::ffi::Tuple<uint32_t, uint32_t>{old_num_cta, old_block_size};
+  }
+
+ protected:
+  AllReduceData* allocate_graph_capture_input(void* data_ptr) {
+    const auto count = registered_count();
+    RuntimeCheck(count < m_graph_buffer_count, "Graph buffer overflow, increase `graph_buffer_count`!");
+    m_graph_capture_inputs.push_back(data_ptr);
+    return get_data_ptr(count);
+  }
+  AllReduceData* get_data_ptr(int64_t which = -1) {
+    const auto count = registered_count();
+    RuntimeCheck(which >= -1 && which < count, "Invalid graph buffer index: ", which, ", count: ", count);
+    const auto start = get_pull_params(m_storage);
+    return static_cast<AllReduceData*>(start) + (1 + which);
+  }
+  int64_t registered_count() const {
+    return static_cast<int64_t>(m_graph_capture_inputs.size());
+  }
+  int64_t pull_signal_bytes() const {
+    return sizeof(device::distributed::Semaphore) * m_max_num_cta_pull;
+  }
+  int64_t push_signal_bytes() const {
+    return sizeof(device::distributed::Semaphore) * m_max_num_cta_push;
+  }
+  int64_t params_bytes() const {
+    return sizeof(AllReduceData) * (1 + m_graph_buffer_count);  // 1 for default
+  }
+  int64_t push_all_ranks_bytes() const {
+    return PushController::kNumStages * m_num_gpu * m_push_buffer_bytes;
+  }
+  int64_t storage_bytes() const {
+    // | SignalArray (pull + push) | GraphBuffers (pull params) | Buffers (pull + push) |
+    return _get_offset_impl(5);
+  }
+  void* get_pull_signal(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(0));
+  }
+  void* get_push_signal(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(1));
+  }
+  void* get_pull_params(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(2));
+  }
+  void* get_pull_buffer(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(3));
+  }
+  void* get_push_buffer(void* ptr) const {
+    return pointer::offset(ptr, _get_offset_impl(4));
+  }
+  int64_t _get_offset_impl(int64_t which) const {
+    const int64_t offset_map[5] = {
+        /*[0]=*/pull_signal_bytes(),
+        /*[1]=*/push_signal_bytes(),
+        /*[2]=*/params_bytes(),
+        /*[3]=*/m_pull_buffer_bytes,
+        /*[4]=*/push_all_ranks_bytes(),
+    };
+    RuntimeCheck(which >= 0 && which <= 5, "Invalid offset index: ", which);
+    return std::accumulate(offset_map, offset_map + which, int64_t(0));
+  }
+
+  const int64_t m_pull_buffer_bytes;
+  const int64_t m_push_buffer_bytes;
+  const int64_t m_graph_buffer_count;
+  const uint32_t m_rank;
+  const uint32_t m_num_gpu;
+  const uint32_t m_max_num_cta_pull;
+  const uint32_t m_max_num_cta_push;
+  // these 2 config should only affect pull kernel
+  uint32_t m_num_cta;
+  uint32_t m_cta_size;
+  // other states
+  bool m_is_graph_capturing = false;
+  int64_t m_cum_registered_count = 0;
+  std::optional<PullController> m_pull_ctrl;
+  std::optional<PushController> m_push_ctrl;
+  void* m_storage = nullptr;
+  std::vector<void*> m_graph_capture_inputs;
+  std::vector<void*> m_peer_storage;
+  std::unordered_map<cudaIpcMemHandle_t, void*, HandleHash, HandleEqual> m_ipc_cache;
+};
+
+struct CustomAllReduceRef : public tvm::ffi::ObjectRef {
+  TVM_FFI_DEFINE_OBJECT_REF_METHODS_NOTNULLABLE(CustomAllReduceRef, tvm::ffi::ObjectRef, CustomAllReduceBase);
+};
+
+}  // namespace host::distributed
+
+namespace device::distributed {
+
+template <typename DType2, size_t N, uint32_t M>
+SGL_DEVICE auto reduce_impl(AlignedVector<DType2, N> (&storage)[M]) -> AlignedVector<DType2, N> {
+  fp32x2_t acc[N] = {};
+#pragma unroll  // unroll num gpu
+  for (uint32_t i = 0; i < M; ++i) {
+#pragma unroll  // unroll vec
+    for (uint32_t j = 0; j < N; ++j) {
+      const auto [x, y] = cast<fp32x2_t>(storage[i][j]);
+      auto& [x_acc, y_acc] = acc[j];
+      x_acc += x;
+      y_acc += y;
+    }
+  }
+
+  AlignedVector<DType2, N> result;
+#pragma unroll
+  for (uint32_t j = 0; j < N; ++j) {
+    result[j] = cast<DType2>(acc[j]);
+  }
+
+  return result;
+}
+
+}  // namespace device::distributed
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/ffi.h b/python/sglang/jit_kernel/include/sgl_kernel/ffi.h
new file mode 100644
index 000000000..17d9048d4
--- /dev/null
+++ b/python/sglang/jit_kernel/include/sgl_kernel/ffi.h
@@ -0,0 +1,104 @@
+#pragma once
+#include <sgl_kernel/utils.h>
+
+#include <dlpack/dlpack.h>
+#include <tvm/ffi/container/shape.h>
+#include <tvm/ffi/container/tensor.h>
+#include <tvm/ffi/extra/c_env_api.h>
+
+#include <algorithm>
+#include <cstdint>
+#include <cstdlib>
+#include <memory>
+#include <optional>
+
+namespace host::ffi {
+
+using tvm::ffi::Tensor, tvm::ffi::TensorView, tvm::ffi::ShapeView;
+
+inline Tensor empty(ShapeView shape, DLDataType dtype, DLDevice device) {
+  return Tensor::FromEnvAlloc(::TVMFFIEnvTensorAlloc, shape, dtype, device);
+}
+
+inline Tensor empty_like(TensorView tensor) {
+  return empty(tensor.shape(), tensor.dtype(), tensor.device());
+}
+
+struct _dummy_deleter {
+  void operator()(void*) const {}
+};
+
+// template <typename Fn = _dummy_deleter>
+
+template <typename Fn>
+struct FromBlobContext {
+  [[no_unique_address]] Fn deleter;
+  int64_t dimension;
+  int64_t* get_shape() {
+    return reinterpret_cast<int64_t*>(this + 1);
+  }
+  int64_t* get_stride() {
+    return this->get_shape() + dimension;
+  }
+};
+
+template <typename Fn = _dummy_deleter>
+inline Tensor from_blob(
+    void* data,
+    ShapeView shape,
+    DLDataType dtype,
+    DLDevice device,
+    Fn&& deleter = {},
+    std::optional<ShapeView> stride = {},
+    uint64_t byte_offset = 0) {
+  using Context = FromBlobContext<std::decay_t<Fn>>;
+  const auto ndim = shape.size();
+  const auto ctx = [&] {
+    auto ptr = std::malloc(sizeof(Context) + sizeof(int64_t) * ndim * 2);
+    auto ctx = static_cast<Context*>(ptr);
+    std::construct_at(ctx, std::forward<Fn>(deleter), static_cast<int64_t>(ndim));
+    stdr::copy_n(shape.data(), ndim, ctx->get_shape());
+    if (stride.has_value()) {
+      RuntimeCheck(stride->size() == ndim, "Stride ndim mismatch!");
+      stdr::copy_n(stride->data(), ndim, ctx->get_stride());
+    } else {
+      int64_t stride_val = 1;
+      for (const auto i : irange(ndim)) {
+        const auto j = ndim - 1 - i;
+        ctx->get_stride()[j] = stride_val;
+        stride_val *= shape[j];
+      }
+    }
+    return ctx;
+  }();
+  const auto tensor = DLTensor{
+      .data = data,
+      .device = device,
+      .ndim = static_cast<int32_t>(ndim),
+      .dtype = dtype,
+      .shape = ctx->get_shape(),
+      .strides = ctx->get_stride(),
+      .byte_offset = byte_offset,
+  };
+  const auto blob_deleter = [](DLManagedTensor* self) {
+    auto ctx = static_cast<Context*>(self->manager_ctx);
+    ctx->deleter(self->dl_tensor.data);
+    std::destroy_at(ctx);
+    std::free(ctx);
+  };
+  auto managed_tensor = DLManagedTensor{tensor, ctx, blob_deleter};
+  return Tensor::FromDLPack(&managed_tensor);
+}
+
+template <typename Fn = _dummy_deleter>
+inline Tensor from_blob_like(
+    void* data,
+    TensorView t,
+    Fn&& deleter = {},
+    bool is_contiguous = false,  // if override to true, the stride will be ignored
+    uint64_t byte_offset = 0) {
+  const auto stride = is_contiguous ? std::nullopt : std::optional{t.strides()};
+  return from_blob(data, t.shape(), t.dtype(), t.device(), std::forward<Fn>(deleter), stride, byte_offset);
+}
+
+}  // namespace host::ffi
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
index 48cb52cf3..b255f12a9 100644
--- a/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
+++ b/python/sglang/jit_kernel/include/sgl_kernel/utils.cuh
@@ -142,6 +142,11 @@ SGL_DEVICE void PDLTriggerSecondary() {
 #endif
 }
 
+template <std::integral T, std::integral U>
+SGL_DEVICE constexpr auto div_ceil(T a, U b) {
+  return (a + b - 1) / b;
+}
+
 /**
  * \brief Load data with the specified type and offset from a void pointer.
  * \tparam T The type to load.
diff --git a/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh b/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
index 86b80c789..fe91325ee 100644
--- a/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
+++ b/python/sglang/jit_kernel/include/sgl_kernel/vec.cuh
@@ -80,15 +80,11 @@ struct AlignedVector {
 
  public:
   /// \brief Vectorized load from `ptr` at the given element `offset`.
-  template <typename U>
-  SGL_DEVICE void load(const U* ptr, std::size_t offset = 0) {
-    static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
+  SGL_DEVICE void load(const void* ptr, int64_t offset = 0) {
     m_storage = reinterpret_cast<const storage_t*>(ptr)[offset];
   }
   /// \brief Vectorized store to `ptr` at the given element `offset`.
-  template <typename U>
-  SGL_DEVICE void store(U* ptr, std::size_t offset = 0) const {
-    static_assert(std::is_same_v<U, T> || std::is_same_v<U, void>);
+  SGL_DEVICE void store(void* ptr, int64_t offset = 0) const {
     reinterpret_cast<storage_t*>(ptr)[offset] = m_storage;
   }
   /// \brief Fill all N elements with the same `value`.
diff --git a/python/sglang/jit_kernel/tests/test_custom_all_reduce.py b/python/sglang/jit_kernel/tests/test_custom_all_reduce.py
new file mode 100644
index 000000000..e1f08dbb9
--- /dev/null
+++ b/python/sglang/jit_kernel/tests/test_custom_all_reduce.py
@@ -0,0 +1,227 @@
+"""
+Correctness test for the JIT custom all-reduce (v2) kernel.
+
+The test compares the JIT custom all-reduce output against NCCL all-reduce
+for various tensor sizes and dtypes, in both eager and CUDA-graph modes.
+
+Usage:
+    python -m pytest test_jit_custom_all_reduce.py -v
+
+This file doubles as the torchrun worker script.  The test class launches
+    torchrun --nproc_per_node=N <this_file>
+and asserts that all worker processes exit successfully.
+"""
+
+from __future__ import annotations
+
+import itertools
+import logging
+import os
+import subprocess
+from typing import Optional
+
+import pytest
+import torch
+import torch.distributed as dist
+from tqdm import tqdm
+
+import sglang.srt.distributed.parallel_state as ps
+from sglang.jit_kernel.all_reduce import AllReduceAlgo
+from sglang.srt.distributed.device_communicators.custom_all_reduce_v2 import (
+    CustomAllReduceV2,
+)
+
+# ---------------------------------------------------------------------------
+# Test parameters (shared between test class and worker)
+# ---------------------------------------------------------------------------
+
+TEST_SIZES = [
+    16,
+    32,
+    512,
+    1024,
+    1024 + 16,  # weird case
+    4 * 1024,
+    32 * 1024,
+    256 * 1024,
+    2 * 1024 * 1024,  # 2M elements
+    4 * 1024 * 1024,  # 4M elements
+]
+TEST_DTYPES = [torch.float16, torch.bfloat16, torch.float32]
+SHOTS = [
+    AllReduceAlgo.ONE_SHOT_PULL,
+    AllReduceAlgo.ONE_SHOT_PUSH,
+    AllReduceAlgo.TWO_SHOT_PULL,
+]
+USE_GRAPH_OPTIONS = [True, False]
+TEST_CONFIG = itertools.product(TEST_SIZES, TEST_DTYPES, SHOTS, USE_GRAPH_OPTIONS)
+TEST_LAYERS = 2
+TEST_LOOP = 16
+
+# ---------------------------------------------------------------------------
+# Test class (runs via pytest, launches torchrun subprocesses)
+# ---------------------------------------------------------------------------
+
+
+def run_torchrun(nproc: int, timeout: int = 300) -> None:
+    """Launch this script as a torchrun worker and assert success."""
+    cmd = [
+        "torchrun",
+        f"--nproc_per_node={nproc}",
+        __file__,
+    ]
+    os.environ["DISABLE_PBAR"] = "1"
+    result = subprocess.run(
+        cmd,
+        stdout=subprocess.PIPE,
+        stderr=subprocess.STDOUT,
+        text=True,
+        timeout=timeout,
+    )
+    assert result.returncode == 0, (
+        f"torchrun (nproc={nproc}) failed with rc={result.returncode}\n"
+        f"{result.stdout}"
+    )
+
+
+@pytest.mark.parametrize("nproc", [2, 3, 4, 5, 6, 7, 8])
+def test_custom_allreduce(nproc: int) -> None:
+    device_count = torch.cuda.device_count()
+    if device_count < nproc:
+        pytest.skip(
+            f"Requires at least {nproc} GPUs, but only {device_count} available"
+        )
+    run_torchrun(nproc)
+
+
+# ---------------------------------------------------------------------------
+# Worker logic (executed by each torchrun process)
+# ---------------------------------------------------------------------------
+
+
+def init_distributed():
+    """Initialize distributed groups via torchrun env vars.
+
+    Returns (rank, device, cpu_group, nccl_group, comm).
+    """
+    local_rank = int(os.environ["LOCAL_RANK"])
+    world_size = int(os.environ["WORLD_SIZE"])
+    rank = local_rank
+    device = torch.device(f"cuda:{rank}")
+    torch.cuda.set_device(device)
+
+    dist.init_process_group(backend="gloo")
+    ps._WORLD = coord = ps.init_world_group(
+        ranks=list(range(world_size)),
+        local_rank=local_rank,
+        backend="nccl",
+    )
+
+    cpu_group = coord.cpu_group
+    nccl_group = coord.device_group
+    assert nccl_group is not None
+
+    max_size = max(TEST_SIZES) * 4
+    comm = CustomAllReduceV2(cpu_group, device, max_size, max_size)
+    if comm.disabled:
+        raise RuntimeError("JIT CustomAllReduceV2 is disabled on this system")
+
+    return rank, device, cpu_group, nccl_group, comm
+
+
+@torch.inference_mode()
+def worker_test(
+    device: torch.device,
+    nccl_group: dist.ProcessGroup,
+    comm: CustomAllReduceV2,
+    size: int,
+    dtype: torch.dtype,
+    use_graph: bool,
+    algo: AllReduceAlgo,
+) -> Optional[RuntimeError]:
+    comm.override_algo = algo
+
+    def get_run_graph_fn():
+        graph = torch.cuda.CUDAGraph()
+        graph_inp = torch.zeros((TEST_LAYERS, size), dtype=dtype, device=device)
+        out_jits = []
+        with comm.capture():
+            with torch.cuda.graph(graph):
+                for i in range(TEST_LAYERS):
+                    out_jits.append(comm.custom_all_reduce(graph_inp[i]))
+                out_jit = torch.stack(out_jits)
+        torch.cuda.synchronize()
+
+        def run_graph(x: torch.Tensor) -> torch.Tensor:
+            graph_inp.copy_(x)
+            graph.replay()
+            return out_jit.clone()
+
+        return run_graph
+
+    def get_run_eager_fn():
+        def run_eager(x: torch.Tensor) -> torch.Tensor:
+            eager_inp = x.clone()
+            out_eagers = []
+            for i in range(TEST_LAYERS):
+                out_eagers.append(comm.custom_all_reduce(eager_inp[i]))
+                torch.cuda.synchronize()
+            return torch.stack(out_eagers)
+
+        return run_eager
+
+    run_fn = get_run_graph_fn() if use_graph else get_run_eager_fn()
+    num_errors = 0
+    for _ in range(TEST_LOOP):
+        # NOTE: 15 * 8 < 128, which is the precision limit for bf16
+        inp = torch.randint(0, 16, (TEST_LAYERS, size), dtype=dtype, device=device)
+        assert comm.should_custom_ar(inp[0])
+        out_ref = inp.clone()
+        dist.all_reduce(out_ref, group=nccl_group)
+        out_jit = run_fn(inp)
+        num_errors += not torch.all(out_jit == out_ref)
+        torch.cuda.synchronize()
+        nccl_group.barrier().wait()
+    if num_errors > 0:
+        return RuntimeError(
+            f"Test failed for {size=}, {dtype=}, {algo=}, "
+            f"{use_graph=} with {num_errors} errors. "
+        )
+    return None
+
+
+def worker_main() -> None:
+    """Entry point for each torchrun worker process."""
+    rank, device, cpu_group, nccl_group, comm = init_distributed()
+    world_size = dist.get_world_size()
+
+    torch.cuda.set_stream(torch.cuda.Stream())
+
+    logging.disable(logging.INFO)  # Suppress internal logging for cleaner test output
+    items = list(enumerate(TEST_CONFIG))
+    disable_pbar = os.environ.get("DISABLE_PBAR", "0") == "1" or rank != 0
+    pbar = tqdm(items, desc=f"Testing {world_size} GPUs", disable=disable_pbar)
+    for i, (size, dtype, algo, use_graph) in pbar:
+        error = worker_test(device, nccl_group, comm, size, dtype, use_graph, algo)
+        if error is not None:
+            print(
+                f"Worker {rank} failed for {size=}, {dtype=}, "
+                f"{algo=}, {use_graph=}, iteration={i}\n"
+                f"Error: {error}"
+            )
+        # communicate the result to rank 0 for logging
+        result = torch.tensor([int(error is not None)], device=device)
+        dist.all_reduce(result, group=cpu_group)
+        failed = bool(result.item())
+        if failed:
+            raise RuntimeError(
+                f"Test failed on rank {rank} for config: "
+                f"{size=}, {dtype=}, {algo=}, {use_graph=}"
+            )
+
+    comm.close()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    worker_main()
diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
index 5e852a080..d7e0a3146 100644
--- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
+++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce.py
@@ -2,7 +2,6 @@
 
 import ctypes
 import logging
-import os
 from contextlib import contextmanager
 from functools import partial
 from typing import Any, List, Optional, Union
@@ -15,11 +14,9 @@ import sglang.srt.distributed.device_communicators.custom_all_reduce_ops as ops
 from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
 from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
-    gpu_p2p_access_check,
-    is_full_nvlink,
+    can_use_custom_all_reduce_with_nvlink,
     is_weak_contiguous,
 )
-from sglang.srt.distributed.parallel_state import in_the_same_node_as
 from sglang.srt.environ import envs
 from sglang.srt.utils import (
     get_bool_env_var,
@@ -36,20 +33,6 @@ _is_musa = is_musa()
 logger = logging.getLogger(__name__)
 
 
-def _can_p2p(rank: int, world_size: int) -> bool:
-    # SGLANG_SKIP_P2P_CHECK can be set to False in sglang
-    SGLANG_SKIP_P2P_CHECK = os.getenv("SGLANG_SKIP_P2P_CHECK", "0") == "1"
-    for i in range(world_size):
-        if i == rank:
-            continue
-        if SGLANG_SKIP_P2P_CHECK:
-            logger.info("Skipping P2P check and trusting the driver's P2P report.")
-            return torch.cuda.can_device_access_peer(rank, i)
-        if not gpu_p2p_access_check(rank, i):
-            return False
-    return True
-
-
 class CustomAllreduce:
     _SUPPORTED_WORLD_SIZES = [2, 4, 6, 8]
     _MAX_CAR_SIZE = 8192 * 1024
@@ -86,35 +69,8 @@ class CustomAllreduce:
             # e.g. in a non-cuda environment
             return
 
-        self.group = group
-
-        assert (
-            dist.get_backend(group) != dist.Backend.NCCL
-        ), "CustomAllreduce should be attached to a non-NCCL group."
-
-        if not all(in_the_same_node_as(group, source_rank=0)):
-            # No need to initialize custom allreduce for multi-node case.
-            logger.warning(
-                "Custom allreduce is disabled because this process group"
-                " spans across nodes."
-            )
-            return
-
-        rank = dist.get_rank(group=self.group)
-        world_size = dist.get_world_size(group=self.group)
-        if world_size == 1:
-            # No need to initialize custom allreduce for single GPU case.
-            return
-
-        if world_size not in CustomAllreduce._SUPPORTED_WORLD_SIZES:
-            logger.warning(
-                "Custom allreduce is disabled due to an unsupported world"
-                " size: %d. Supported world sizes: %s. To silence this "
-                "warning, specify disable_custom_all_reduce=True explicitly.",
-                world_size,
-                str(CustomAllreduce._SUPPORTED_WORLD_SIZES),
-            )
-            return
+        rank = dist.get_rank(group=group)
+        world_size = dist.get_world_size(group=group)
 
         if isinstance(device, int):
             device = torch.device(f"cuda:{device}")
@@ -123,46 +79,16 @@ class CustomAllreduce:
         # now `device` is a `torch.device` object
         assert isinstance(device, torch.device)
         self.device = device
+        full_nvlink = can_use_custom_all_reduce_with_nvlink(
+            group=group,
+            device=device,
+            supported_world_size=self._SUPPORTED_WORLD_SIZES,
+            cls_name="CustomAllreduce",
+        )
+        if full_nvlink is None:
+            return  # fail to get nvlink status
 
-        cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
-        if cuda_visible_devices:
-            device_ids = list(map(int, cuda_visible_devices.split(",")))
-        else:
-            device_ids = list(range(torch.cuda.device_count()))
-
-        physical_device_id = device_ids[device.index]
-        tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
-        gather_list = [
-            torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
-        ]
-        dist.all_gather(gather_list, tensor, group=self.group)
-        physical_device_ids = [t.item() for t in gather_list]
-
-        # test nvlink first, this will filter out most of the cases
-        # where custom allreduce is not supported
-        # this checks hardware and driver support for NVLink
-        if _is_cuda or _is_hip or _is_musa:
-            full_nvlink = is_full_nvlink(physical_device_ids, world_size)
-
-        if world_size > 2 and not full_nvlink:
-            logger.warning(
-                "Custom allreduce is disabled because it's not supported on"
-                " more than two PCIe-only GPUs. To silence this warning, "
-                "specify disable_custom_all_reduce=True explicitly."
-            )
-            return
-        # test P2P capability, this checks software/cudaruntime support
-        # this is expensive to compute at the first time
-        # then we cache the result
-        # On AMD GPU, p2p is always enabled between XGMI connected GPUs
-        if not _is_hip and not _can_p2p(rank, world_size):
-            logger.warning(
-                "Custom allreduce is disabled because your platform lacks "
-                "GPU P2P capability or P2P test failed. To silence this "
-                "warning, specify disable_custom_all_reduce=True explicitly."
-            )
-            return
-
+        self.group = group
         self.max_size = max_size
         self.rank = rank
         self.world_size = world_size
@@ -458,7 +384,17 @@ def dispatch_custom_allreduce():
 
     On AMD with 1-stage AR enabled, use sglang's CustomAllreduce (has deterministic_all_reduce method).
     Otherwise use AiterCustomAllreduce if available.
+
+    Set SGLANG_USE_JIT_ALL_REDUCE=1 to use the JIT-compiled v2 implementation.
     """
+    # HARDCODED: opt-in flag for v2 JIT all-reduce.
+    # Set SGLANG_USE_JIT_ALL_REDUCE=1 to enable.
+    if _is_cuda and get_bool_env_var("SGLANG_USE_JIT_ALL_REDUCE", default="false"):
+        from .custom_all_reduce_v2 import CustomAllReduceV2
+
+        logger.debug("[AR] Using CustomAllReduceV2 (JIT-compiled)")
+        return CustomAllReduceV2
+
     if _is_cuda or _is_musa:
         return CustomAllreduce
 
diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py
index c74cd58d6..b1505177a 100644
--- a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py
+++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_utils.py
@@ -18,6 +18,7 @@ import torch.multiprocessing as mp
 from typing_extensions import ParamSpec
 
 from sglang.srt.distributed.device_communicators.cuda_wrapper import CudaRTLibrary
+from sglang.srt.distributed.parallel_state import in_the_same_node_as
 from sglang.srt.utils import is_cuda, is_hip, is_musa
 
 logger = logging.getLogger(__name__)
@@ -384,7 +385,92 @@ def is_weak_contiguous(inp: torch.Tensor):
     )
 
 
-__all__ = ["gpu_p2p_access_check"]
+def can_p2p(rank: int, world_size: int) -> bool:
+    # SGLANG_SKIP_P2P_CHECK can be set to False in sglang
+    SGLANG_SKIP_P2P_CHECK = os.getenv("SGLANG_SKIP_P2P_CHECK", "0") == "1"
+    for i in range(world_size):
+        if i == rank:
+            continue
+        if SGLANG_SKIP_P2P_CHECK:
+            logger.info("Skipping P2P check and trusting the driver's P2P report.")
+            return torch.cuda.can_device_access_peer(rank, i)
+        if not gpu_p2p_access_check(rank, i):
+            return False
+    return True
+
+
+def can_use_custom_all_reduce_with_nvlink(
+    group: torch.distributed.ProcessGroup,
+    device: torch.device,
+    supported_world_size: List[int],
+    cls_name: str,
+) -> Optional[bool]:  # None if fail; otherwise return whether NVLink is available
+    assert (
+        dist.get_backend(group) != dist.Backend.NCCL
+    ), f"{cls_name} should be attached to a non-NCCL group."
+
+    rank = dist.get_rank(group=group)
+    world_size = dist.get_world_size(group=group)
+
+    # No need to initialize custom allreduce for single GPU case.
+    if world_size == 1:
+        return
+
+    # No need to initialize custom allreduce for multi-node case.
+    if not all(in_the_same_node_as(group, source_rank=0)):
+        logger.warning(
+            f"{cls_name} is disabled because this process group" " spans across nodes."
+        )
+        return
+
+    # For not supported world size, we disable custom allreduce.
+    if world_size not in supported_world_size:
+        logger.warning(
+            f"{cls_name} is disabled due to an unsupported world"
+            f" size: {world_size}. Supported world sizes: {supported_world_size}. "
+            "To silence this warning, specify disable_custom_all_reduce=True explicitly.",
+        )
+        return
+
+    cuda_visible_devices = os.environ.get("CUDA_VISIBLE_DEVICES", None)
+    if cuda_visible_devices:
+        device_ids = list(map(int, cuda_visible_devices.split(",")))
+    else:
+        device_ids = list(range(torch.cuda.device_count()))
+    physical_device_id = device_ids[device.index]
+    tensor = torch.tensor([physical_device_id], dtype=torch.int, device="cpu")
+    gather_list = [
+        torch.tensor([0], dtype=torch.int, device="cpu") for _ in range(world_size)
+    ]
+    dist.all_gather(gather_list, tensor, group=group)
+    physical_device_ids = [int(t) for t in gather_list]
+    full_nvlink = is_full_nvlink(physical_device_ids, world_size)
+
+    # test nvlink first, this will filter out most of the cases
+    # where custom allreduce is not supported
+    # this checks hardware and driver support for NVLink
+    if world_size > 2 and not full_nvlink:
+        logger.warning(
+            f"{cls_name} is disabled because it's not supported on"
+            " more than two PCIe-only GPUs. To silence this warning, "
+            "specify disable_custom_all_reduce=True explicitly."
+        )
+        return
+
+    # test P2P capability, this checks software/cudaruntime support
+    # this is expensive to compute at the first time
+    # then we cache the result
+    # On AMD GPU, p2p is always enabled between XGMI connected GPUs
+    if not _is_hip and not can_p2p(rank, world_size):
+        logger.warning(
+            f"{cls_name} is disabled because your platform lacks "
+            "GPU P2P capability or P2P test failed. To silence this "
+            "warning, specify disable_custom_all_reduce=True explicitly."
+        )
+        return
+
+    return full_nvlink
+
 
 if __name__ == "__main__":
     batch_src, batch_tgt, output_file = pickle.loads(sys.stdin.buffer.read())
diff --git a/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
new file mode 100644
index 000000000..2006554f9
--- /dev/null
+++ b/python/sglang/srt/distributed/device_communicators/custom_all_reduce_v2.py
@@ -0,0 +1,189 @@
+import logging
+from contextlib import contextmanager
+from dataclasses import dataclass, replace
+from typing import Dict, List, Optional, TypeVar
+
+import torch
+import torch.distributed as dist
+from torch.distributed import ProcessGroup
+
+from sglang.jit_kernel.all_reduce import AllReduceAlgo, get_custom_all_reduce_cls
+from sglang.srt.distributed import is_in_piecewise_cuda_graph
+from sglang.srt.distributed.device_communicators.custom_all_reduce_utils import (
+    can_use_custom_all_reduce_with_nvlink,
+    is_weak_contiguous,
+)
+from sglang.srt.utils import is_sm100_supported, log_info_on_rank0
+
+logger = logging.getLogger(__name__)
+
+T = TypeVar("T")
+
+INF = 1 << 60
+
+
+@dataclass(frozen=True)
+class ModeConfig:
+    one_shot_push_threshold: int  # below this, use one-shot push
+    one_shot_pull_threshold: int  # below this, use one-shot pull
+
+
+class CustomAllReduceV2:
+    def __init__(
+        self,
+        group: ProcessGroup,
+        device: torch.device,
+        max_pull_size: Optional[int] = None,
+        max_push_size: Optional[int] = None,
+    ) -> None:
+        _init_config()
+        self.disabled = True
+        full_nvlink = can_use_custom_all_reduce_with_nvlink(
+            group=group,
+            device=device,
+            supported_world_size=list(THRESHOLD_2_SHOT_MAP.keys()),
+            cls_name="CustomAllReduceV2",
+        )
+        if full_nvlink != True:
+            return
+
+        self.group = group
+        self.rank = dist.get_rank(group=self.group)
+        self.world_size = dist.get_world_size(group=self.group)
+        self.override_shot(None)
+        if max_pull_size is None:
+            max_pull_size = 16 * 1024 * 1024  # default to 16MB
+        if max_push_size is None:
+            max_push_size = self.config.one_shot_push_threshold
+        max_push_size = min(max_push_size, max_pull_size)
+        self.max_pull_size = max_pull_size
+        self.max_push_size = max_push_size
+        self.override_algo: Optional[AllReduceAlgo] = None
+        self.obj = get_custom_all_reduce_cls()(
+            rank=self.rank,
+            world_size=self.world_size,
+            pull_buffer_bytes=self.max_pull_size,
+            push_buffer_bytes=self.max_push_size,
+            graph_input_count=131072,
+        )
+        self._post_init_obj()
+        self.disabled = False
+        log_info_on_rank0(logger, "Custom allreduce v2 initialized successfully")
+
+    def override_shot(self, shot: int | None):
+        if shot is None:
+            self.config = THRESHOLD_2_SHOT_MAP[self.world_size]
+        else:
+            assert shot in (1, 2)
+            threshold = INF if shot == 1 else 0
+            self.config = replace(self.config, one_shot_pull_threshold=threshold)
+
+    @contextmanager
+    def capture(self):
+        try:
+            self.obj.set_cuda_graph_capture(True)
+            yield
+        finally:
+            self.obj.set_cuda_graph_capture(False)
+        if not self.disabled:
+            # cannot call when graph is capturing
+            assert (
+                torch.cuda.is_current_stream_capturing() == False
+            ), "Cannot register graph inputs while capturing CUDA graph"
+            pairs = self.obj.share_graph_inputs()
+            handles = [handle for _, handle in pairs]
+            offsets = [offset for offset, _ in pairs]
+            handles_all = self._share_list(handles)
+            offsets_all = self._share_list(offsets)
+            result = [list(zip(o, h)) for o, h in zip(offsets_all, handles_all)]
+            self.obj.register_inputs(result)
+            log_info_on_rank0(logger, f"Registering {len(pairs)} cuda graph addresses")
+
+    def should_custom_ar(self, inp: torch.Tensor) -> bool:
+        """Check if the input tensor is suitable for custom all-reduce."""
+        if self.disabled:
+            return False
+        inp_size = inp.numel() * inp.element_size()
+        # custom allreduce requires input byte size to be multiples of 16
+        if inp_size % 16 != 0:
+            return False
+        if not is_weak_contiguous(inp):
+            return False
+        return inp_size <= self.max_pull_size
+
+    def custom_all_reduce(self, input: torch.Tensor) -> torch.Tensor:
+        if is_in_piecewise_cuda_graph():  # disable inplace optimization
+            try:
+                self.obj.set_cuda_graph_capture(False)
+                return self._all_reduce(input)
+            finally:
+                self.obj.set_cuda_graph_capture(True)
+        return self._all_reduce(input)
+
+    def close(self):
+        if not self.disabled and hasattr(self, "obj"):
+            self.obj.free(self.group)
+
+    def _all_reduce(self, input: torch.Tensor) -> torch.Tensor:
+        """Perform the actual all-reduce via JIT kernel."""
+        algo = self._determine_algo(input)
+        return torch.from_dlpack(self.obj.all_reduce(input, algo))
+
+    def _determine_algo(self, input: torch.Tensor) -> AllReduceAlgo:
+        if self.override_algo is not None:
+            return self.override_algo
+        input_bytes = input.numel() * input.element_size()
+        if input_bytes <= self.config.one_shot_push_threshold:
+            return AllReduceAlgo.ONE_SHOT_PUSH
+        if input_bytes <= self.config.one_shot_pull_threshold:
+            return AllReduceAlgo.ONE_SHOT_PULL
+        else:
+            return AllReduceAlgo.TWO_SHOT_PULL
+
+    def _post_init_obj(self):
+        handles = [self.obj.share_storage()]
+        result = self._share_list(handles)
+        assert all(len(r) == 1 for r in result)
+        result = [h[0] for h in result]
+        self.obj.post_init(result)
+
+    def _share_list(self, input: List[T]) -> List[List[T]]:
+        input_tensor = torch.tensor(input, dtype=torch.int64, device="cpu")
+        gather_list = [torch.empty_like(input_tensor) for _ in range(self.world_size)]
+        dist.all_gather(gather_list, input_tensor, group=self.group)
+        return [g.tolist() for g in gather_list]
+
+    def __del__(self):
+        self.close()
+
+
+def _init_config():
+    global THRESHOLD_2_SHOT_MAP
+    KB, MB = 1024, 1024 * 1024
+
+    if is_sm100_supported():
+        # NOTE: This result is based on benchmarks on B200 GPUs
+        THRESHOLD_2_SHOT_MAP = {
+            2: ModeConfig(4 * MB, INF),
+            3: ModeConfig(4 * MB, 4 * MB),
+            4: ModeConfig(2 * MB, 2 * MB),
+            5: ModeConfig(2 * MB, 2 * MB),
+            6: ModeConfig(1 * MB, 1 * MB),
+            7: ModeConfig(896 * KB, 896 * KB),
+            8: ModeConfig(720 * KB, 720 * KB),
+        }
+    else:
+        # NOTE: This result is based on benchmarks on H200 GPUs
+        THRESHOLD_2_SHOT_MAP = {
+            2: ModeConfig(2 * MB, INF),
+            3: ModeConfig(512 * KB, 512 * KB),
+            4: ModeConfig(384 * KB, 256 * KB),
+            5: ModeConfig(256 * KB, 256 * KB),
+            6: ModeConfig(192 * KB, 192 * KB),
+            7: ModeConfig(192 * KB, 192 * KB),
+            8: ModeConfig(160 * KB, 160 * KB),
+        }
+    # TODO: tune on more GPUs, e.g A100
+
+
+THRESHOLD_2_SHOT_MAP: Dict[int, ModeConfig] = {}