fix benchmark

Signed-off-by: Woosuk Kwon <woosuk@inferact.ai>

csrc/standalone/custom_ar/bench_moe_allgather.py

@@ -0,0 +1,200 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Benchmark: Lamport all-gather vs NCCL."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def gpu_timer(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    """Time with CUDA graph to exclude CPU overhead."""
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import _load_lib
+
+        lib = _load_lib()
+    dist.barrier()
+    if rank != 0:
+        from moe_allgather import _load_lib
+
+        lib = _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    configs = [
+        ("1tok", 1),
+        ("2tok", 2),
+        ("4tok", 4),
+        ("8tok", 8),
+        ("16tok", 16),
+        ("32tok", 32),
+        ("64tok", 64),
+        ("128tok", 128),
+        ("256tok", 256),
+    ]
+    topk = 8
+    hd = 3584
+    sd = 448
+
+    if rank == 0:
+        print(f"world_size={ws}, max_per_rank={ag.max_per_rank} bytes")
+        print(f"{'config':<12} {'lamport_graph':>10} {'nccl_graph':>10} {'speedup':>8}")
+        print("-" * 65)
+
+    for name, N in configs:
+        # Check if data fits in buffer.
+        cursor = 0
+        per_tok = topk * 4 + topk * 4 + hd + sd
+        cursor = N * per_tok
+        cursor = (cursor + 15) & ~15
+        if cursor > ag.max_per_rank:
+            if rank == 0:
+                print(f"{name:<12} {'skip (too large)':>40}")
+            continue
+
+        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
+        wt = torch.randn(N, topk, dtype=torch.float32, device=dev).abs()
+        hs = torch.randint(0, 255, (N, hd), dtype=torch.uint8, device=dev)
+        sc = torch.randint(0, 255, (N, sd), dtype=torch.uint8, device=dev)
+        inputs = [ids, wt, hs, sc]
+
+        # Custom Lamport kernel.
+        c_outs = [
+            torch.empty(N * ws, *t.shape[1:], dtype=t.dtype, device=dev) for t in inputs
+        ]
+
+        def run_lamport():
+            lib.moe_all_gather(
+                ag._buf_ptrs_ptr,
+                ag._counters_ptr,
+                rank,
+                ws,
+                ag.seg_capacity,
+                ag.rank_stride,
+                inputs,
+                c_outs,
+            )
+
+        # Lamport with CUDA graph.
+        try:
+            lam_g_us = gpu_timer_graph(run_lamport)
+        except Exception as ex:
+            lam_g_us = float("nan")
+            if rank == 0:
+                print(f"  [graph capture failed: {ex}]")
+
+        # NCCL 1×AG (concat into one tensor).
+        cat_inp = torch.cat(
+            [t.reshape(N, -1).contiguous().view(torch.uint8) for t in inputs],
+            dim=1,
+        ).contiguous()
+        cat_out = torch.empty(N * ws, cat_inp.shape[1], dtype=torch.uint8, device=dev)
+
+        def run_nccl():
+            dist.all_gather_into_tensor(cat_out, cat_inp)
+
+        # NCCL with CUDA graph.
+        try:
+            nccl_g_us = gpu_timer_graph(run_nccl)
+        except Exception:
+            nccl_g_us = float("nan")
+
+        if rank == 0:
+            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
+            print(f"{name:<12} {lam_g_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x")
+
+    dist.barrier()
+    dist.destroy_process_group()
+
+
+if __name__ == "__main__":
+    main()

csrc/standalone/custom_ar/moe_allgather.cu

@@ -0,0 +1,301 @@
+// Lamport-based MoE all-gather kernel for EP dispatch.
+//
+// Replaces the flag-barrier approach with a Lamport sentinel protocol
+// (inspired by FlashInfer's trtllm_allreduce_fusion).
+//
+// Key advantages over the flag-barrier approach:
+//   - No explicit barriers (sentinels provide per-element synchronization).
+//   - Push model: NVLink writes (fire-and-forget) instead of NVLink reads.
+//   - Triple buffering: no end barrier needed.
+//
+// Gathers the MoE dispatch tensors from all EP ranks:
+//   - topk_ids      [N, topk]     int32
+//   - topk_weights  [N, topk]     float32 / bfloat16
+//   - hidden_states [N, D_h]      uint8 (NVFP4) / bfloat16
+//   - quant_scales  [N, D_s]      (optional)
+//
+// Double-buffer layout in each rank's IPC buffer:
+//   [Segment 0][Segment 1]
+//   Each segment: [Rank 0 slot][Rank 1 slot]...[Rank N-1 slot]
+//   Each rank slot: packed tensors at 16-byte aligned offsets.
+//
+// Sentinel: 0x80000000 (negative-zero in float32).  The writer replaces
+// any data word matching the sentinel with 0 before pushing.  The reader
+// spin-loads (volatile) until no sentinel words remain in the vector.
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// ---------------------------------------------------------------------------
+// Volatile 128-bit load/store and sentinel helpers
+// ---------------------------------------------------------------------------
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// ---------------------------------------------------------------------------
+// Lamport all-gather kernel
+// ---------------------------------------------------------------------------
+//
+// Phase 1 — PUSH: each rank writes its packed data to ALL peers' current
+//           segment via regular stores (NVLink push, fire-and-forget).
+// Phase 2 — CLEAR: each rank writes sentinels to the OLDEST segment of
+//           its own buffer, preparing it for reuse.
+// Phase 3 — POLL + SCATTER: each rank volatile-loads from its own current
+//           segment, spinning until sentinels disappear, then scatters
+//           directly to per-tensor output arrays.
+// Phase 4 — ADVANCE: one thread advances the triple-buffer ring counter.
+
+template <int ngpus, int nbufs>
+__global__ void __launch_bounds__(512, 1) moe_allgather_lamport_kernel(
+    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
+    int* counters,      // [0] = unused, [1] = ring (0/1/2), [2] = prev total_sz
+    int rank,
+    int seg_capacity,  // bytes per segment
+    int rank_stride,   // bytes per rank-slot within a segment
+    int total_sz,      // int4 units of actual packed data per rank
+    // inputs (up to 4)
+    const void* inp0, const void* inp1, const void* inp2, const void* inp3,
+    int off0, int sz0, int off1, int sz1, int off2, int sz2, int off3, int sz3,
+    // outputs (up to 4)
+    void* out0, void* out1, void* out2, void* out3) {
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  // Read segment index and previous clear size.
+  const int seg = counters[1];            // 0 or 1
+  const int prev_total_sz = counters[2];  // set by previous invocation
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  // Sentinel vector for clearing.
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH local data to ALL peers ----
+  // Write to peer_r's buffer at [rank * rank_stride + off_i].
+
+#define PUSH(idx, inp_ptr, off_val, sz_val)                                 \
+  if constexpr (nbufs > (idx)) {                                            \
+    const V* src = reinterpret_cast<const V*>(inp_ptr);                     \
+    for (int i = tid; i < (sz_val); i += stride) {                          \
+      V val = remove_sentinel(src[i]);                                      \
+      _Pragma("unroll") for (int r = 0; r < ngpus; r++) {                   \
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride + (off_val))[i] = \
+            val;                                                            \
+      }                                                                     \
+    }                                                                       \
+  }
+
+  PUSH(0, inp0, off0, sz0)
+  PUSH(1, inp1, off1, sz1)
+  PUSH(2, inp2, off2, sz2)
+  PUSH(3, inp3, off3, sz3)
+#undef PUSH
+
+  // ---- Phase 2: CLEAR only the previously-written data in oldest segment ----
+  // Only clear what the previous invocation actually wrote (per rank-slot).
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: POLL + SCATTER ----
+  // Volatile-load from own buffer; spin until sentinel gone; scatter to output.
+  char* my = bufs[rank];
+
+#define POLL(idx, out_ptr, off_val, sz_val)                                \
+  if constexpr (nbufs > (idx)) {                                           \
+    for (int i = tid; i < (sz_val); i += stride) {                         \
+      _Pragma("unroll") for (int s = 0; s < ngpus; s++) {                  \
+        V val;                                                             \
+        do {                                                               \
+          val = ld128v(                                                    \
+              reinterpret_cast<V*>(my + s * rank_stride + (off_val)) + i); \
+        } while (has_sentinel(val));                                       \
+        reinterpret_cast<V*>(out_ptr)[s * (sz_val) + i] = val;             \
+      }                                                                    \
+    }                                                                      \
+  }
+
+  POLL(0, out0, off0, sz0)
+  POLL(1, out1, off1, sz1)
+  POLL(2, out2, off2, sz2)
+  POLL(3, out3, off3, sz3)
+#undef POLL
+
+  // ---- Phase 4: ADVANCE ring counter + store clear size for next call ----
+  // Stream serialization ensures the next kernel sees these updates.
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel initialization kernel
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+// ---------------------------------------------------------------------------
+// Host launcher
+// ---------------------------------------------------------------------------
+
+struct TensorDesc {
+  void* inp;
+  int off;
+  int sz;
+  int64_t nbytes;
+};
+
+static TensorDesc make_desc(torch::Tensor& inp, int64_t& cursor) {
+  TORCH_CHECK(inp.is_contiguous(), "input must be contiguous");
+  int64_t nbytes = inp.numel() * inp.element_size();
+  TORCH_CHECK(nbytes % 16 == 0, "tensor byte size must be multiple of 16, got ",
+              nbytes);
+  cursor = (cursor + 15) & ~15;
+  TensorDesc d;
+  d.inp = inp.data_ptr();
+  d.off = static_cast<int>(cursor);
+  d.sz = static_cast<int>(nbytes / 16);
+  d.nbytes = nbytes;
+  cursor += nbytes;
+  return d;
+}
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_all_gather(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
+                    int64_t world_size, int64_t seg_capacity,
+                    int64_t rank_stride, std::vector<torch::Tensor>& inputs,
+                    std::vector<torch::Tensor>& outputs) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(inputs.size());
+  TORCH_CHECK(n >= 2 && n <= 4, "2-4 input tensors required");
+  TORCH_CHECK(inputs.size() == outputs.size());
+
+  int64_t cursor = 0;
+  TensorDesc descs[4] = {};
+  for (int i = 0; i < n; i++) descs[i] = make_desc(inputs[i], cursor);
+  TORCH_CHECK(cursor % 16 == 0);
+  int total_sz = static_cast<int>(cursor / 16);
+  TORCH_CHECK(cursor <= rank_stride, "packed data (", cursor,
+              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
+
+  int ws = static_cast<int>(world_size);
+  for (int i = 0; i < n; i++) {
+    TORCH_CHECK(outputs[i].is_contiguous());
+    TORCH_CHECK(outputs[i].numel() == inputs[i].numel() * ws);
+  }
+
+  int r = static_cast<int>(rank);
+  int threads = 512;
+  int blocks =
+      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  void *inps[4] = {}, *outs[4] = {};
+  int offs[4] = {}, szs[4] = {};
+  for (int i = 0; i < n; i++) {
+    inps[i] = descs[i].inp;
+    offs[i] = descs[i].off;
+    szs[i] = descs[i].sz;
+    outs[i] = outputs[i].data_ptr();
+  }
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+
+#define KL(ng, nb)                                                        \
+  moe_allgather_lamport_kernel<ng, nb><<<blocks, threads, 0, stream>>>(   \
+      bp, ct, r, sc, rs, total_sz, inps[0], inps[1], inps[2], inps[3],    \
+      offs[0], szs[0], offs[1], szs[1], offs[2], szs[2], offs[3], szs[3], \
+      outs[0], outs[1], outs[2], outs[3]);
+
+#define GPU_CASE(ng) \
+  case ng:           \
+    switch (n) {     \
+      case 2:        \
+        KL(ng, 2);   \
+        break;       \
+      case 3:        \
+        KL(ng, 3);   \
+        break;       \
+      case 4:        \
+        KL(ng, 4);   \
+        break;       \
+    }                \
+    break;
+
+  switch (ws) {
+    GPU_CASE(2)
+    GPU_CASE(4)
+    GPU_CASE(6)
+    GPU_CASE(8)
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef GPU_CASE
+#undef KL
+}
+
+// ---------------------------------------------------------------------------
+// Python binding
+// ---------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_all_gather", &moe_all_gather, "Lamport MoE all-gather");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}

csrc/standalone/custom_ar/moe_allgather.py

@@ -0,0 +1,116 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Lamport-based fused MoE all-gather for EP dispatch.
+
+JIT-compiles the CUDA kernel on first use (cached afterwards).
+Uses a Lamport sentinel protocol (push writes + per-element sync)
+with triple-buffered IPC regions — no explicit barriers.
+
+Usage:
+    ag = MoeAllGather(custom_allreduce)
+    ids_g, wt_g, hs_g, sc_g = ag.gather(topk_ids, topk_weights, hidden, scales)
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_allgather.cu"))
+    _lib = load(
+        name="moe_allgather_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_AG_VERBOSE", "") == "1",
+    )
+    return _lib
+
+
+class MoeAllGather:
+    """Lamport-based MoE dispatch all-gather with triple buffering."""
+
+    def __init__(self, ca_comm):
+        self.rank = ca_comm.rank
+        self.world_size = ca_comm.world_size
+        self.device = ca_comm.device
+        self.buffer_ptrs = ca_comm.buffer_ptrs
+        self.max_size = ca_comm.max_size
+
+        ws = self.world_size
+        # Double-buffer layout: 2 segments, each with ws rank-slots.
+        # Safe because kernels in the same stream are serialized, and the
+        # Use the FIRST half of the IPC buffer (second half reserved for
+        # MoeReduceScatter) to avoid overlapping writes.
+        half_size = (self.max_size // 2) & ~15
+
+        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
+        # seg_capacity and rank_stride are 16-byte aligned.
+        self.seg_capacity = (half_size // 2) & ~15
+        self.rank_stride = (self.seg_capacity // ws) & ~15
+        self.max_per_rank = self.rank_stride  # max packed bytes per rank
+
+        # Buffer pointer array on device (no offset — first half).
+        self._buf_ptrs = torch.zeros(
+            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
+        )
+        for i in range(ws):
+            self._buf_ptrs[i] = self.buffer_ptrs[i]
+        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
+
+        # Counters on device: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
+        self._counters = torch.zeros(
+            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
+        )
+        self._counters_ptr = self._counters.data_ptr()
+
+        # Initialize our half with sentinel values.
+        lib = _load_lib()
+        lib.lamport_init(self.buffer_ptrs[self.rank], half_size)
+        torch.accelerator.synchronize(self.device)
+
+    def gather(
+        self,
+        topk_ids: torch.Tensor,
+        topk_weights: torch.Tensor,
+        hidden_states: torch.Tensor,
+        scales: torch.Tensor | None = None,
+    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:
+        lib = _load_lib()
+        ws = self.world_size
+
+        inputs = [topk_ids, topk_weights, hidden_states]
+        if scales is not None:
+            inputs.append(scales)
+
+        outputs = [
+            torch.empty((t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device)
+            for t in inputs
+        ]
+
+        lib.moe_all_gather(
+            self._buf_ptrs_ptr,
+            self._counters_ptr,
+            self.rank,
+            self.world_size,
+            self.seg_capacity,
+            self.rank_stride,
+            inputs,
+            outputs,
+        )
+
+        if scales is not None:
+            return outputs[0], outputs[1], outputs[2], outputs[3]
+        return outputs[0], outputs[1], outputs[2], None

csrc/standalone/custom_ar/moe_reduce_scatter.cu

@@ -0,0 +1,261 @@
+// Lamport-based MoE reduce-scatter kernel for EP combine.
+//
+// JIT-compilable via torch.utils.cpp_extension — no vLLM build required.
+//
+// Reduce-scatters a bf16 tensor [N_total, D] across EP ranks.  Each rank
+// contributes its partial MoE output; the kernel sums all contributions
+// and each rank receives its own slice of the result.
+//
+// Protocol (same as the all-gather variant):
+//   1. PUSH: write own data to all peers' Lamport buffers (NVLink push).
+//   2. CLEAR: write sentinels to old segment of own buffer.
+//   3. POLL + REDUCE: volatile-load all peers' data for own slice,
+//      accumulate in fp32, convert back to bf16, store to output.
+//   4. ADVANCE: toggle double-buffer index.
+//
+// Sentinel: 0x80000000 (two bf16 negative-zeros packed in uint32).
+// For bf16 reduce, replacing -0 with +0 is lossless.
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// ---------------------------------------------------------------------------
+// Volatile 128-bit load and sentinel helpers
+// ---------------------------------------------------------------------------
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// ---------------------------------------------------------------------------
+// bf16 ↔ fp32 helpers for int4 (8 bf16 values = 16 bytes)
+// ---------------------------------------------------------------------------
+
+// Accumulate 8 bf16 values from an int4 into 8 fp32 accumulators.
+static DINLINE void accumulate_bf16(float* acc, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
+}
+
+// Convert 8 fp32 accumulators to bf16 and pack into int4.
+static DINLINE int4 fp32_to_bf16_int4(const float* acc) {
+  int4 out;
+  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
+#pragma unroll
+  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(acc[k]);
+  return out;
+}
+
+// ---------------------------------------------------------------------------
+// Lamport reduce-scatter kernel
+// ---------------------------------------------------------------------------
+
+template <int ngpus>
+__global__ void __launch_bounds__(512, 1) moe_rs_lamport_kernel(
+    int64_t* buf_ptrs,  // [ngpus] IPC buffer base addresses (device)
+    int* counters,      // [0] = unused, [1] = seg (0/1), [2] = prev total_sz
+    int rank,
+    int seg_capacity,   // bytes per segment
+    int rank_stride,    // bytes per rank-slot within a segment
+    const void* input,  // [N_total, D] bf16 — full input
+    void* output,       // [N_per_rank, D] bf16 — this rank's reduced slice
+    int total_sz,       // int4 units of full input per rank
+    int slice_off,      // int4 offset to this rank's slice within packed data
+    int slice_sz) {     // int4 units of this rank's slice
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  // Read segment index and previous clear size.
+  const int seg = counters[1];
+  const int prev_total_sz = counters[2];
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH full input to ALL peers ----
+  {
+    const V* src = reinterpret_cast<const V*>(input);
+    for (int i = tid; i < total_sz; i += stride) {
+      V val = remove_sentinel(src[i]);
+#pragma unroll
+      for (int r = 0; r < ngpus; r++)
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
+    }
+  }
+
+  // ---- Phase 2: CLEAR old segment ----
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: POLL + REDUCE for own slice ----
+  // Read all ranks' data at [slice_off, slice_off + slice_sz) from own buffer,
+  // sum in fp32, store bf16 result.
+  {
+    char* my = bufs[rank];
+    V* dst = reinterpret_cast<V*>(output);
+
+    for (int i = tid; i < slice_sz; i += stride) {
+      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
+
+#pragma unroll
+      for (int s = 0; s < ngpus; s++) {
+        V val;
+        do {
+          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + slice_off +
+                       i);
+        } while (has_sentinel(val));
+        accumulate_bf16(acc, val);
+      }
+
+      dst[i] = fp32_to_bf16_int4(acc);
+    }
+  }
+
+  // ---- Phase 4: ADVANCE ----
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel initialization
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+// ---------------------------------------------------------------------------
+// Host launcher
+// ---------------------------------------------------------------------------
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_reduce_scatter(int64_t buf_ptrs_ptr, int64_t counters_ptr,
+                        int64_t rank, int64_t world_size, int64_t seg_capacity,
+                        int64_t rank_stride, torch::Tensor input,
+                        torch::Tensor output) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  TORCH_CHECK(input.is_contiguous(), "input must be contiguous");
+  TORCH_CHECK(output.is_contiguous(), "output must be contiguous");
+  TORCH_CHECK(input.scalar_type() == torch::kBFloat16,
+              "input must be bf16, got ", input.scalar_type());
+  TORCH_CHECK(output.scalar_type() == torch::kBFloat16, "output must be bf16");
+
+  int ws = static_cast<int>(world_size);
+  int r = static_cast<int>(rank);
+
+  // Input: [N_total, D], Output: [N_per_rank, D]
+  int64_t N_total = input.size(0);
+  int64_t D = input.size(1);
+  TORCH_CHECK(N_total % ws == 0, "N_total must be divisible by world_size");
+  int64_t N_per_rank = N_total / ws;
+  TORCH_CHECK(output.size(0) == N_per_rank);
+  TORCH_CHECK(output.size(1) == D);
+
+  int64_t input_bytes = input.numel() * input.element_size();
+  TORCH_CHECK(input_bytes % 16 == 0,
+              "input byte size must be multiple of 16, got ", input_bytes);
+  TORCH_CHECK(input_bytes <= rank_stride, "input (", input_bytes,
+              " bytes) exceeds rank_stride (", rank_stride, " bytes)");
+
+  int total_sz = static_cast<int>(input_bytes / 16);
+  int slice_sz = total_sz / ws;
+  int slice_off = r * slice_sz;
+
+  int threads = 512;
+  int blocks =
+      std::max(1, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+
+#define LAUNCH(ng)                                                      \
+  moe_rs_lamport_kernel<ng><<<blocks, threads, 0, stream>>>(            \
+      bp, ct, r, sc, rs, input.data_ptr(), output.data_ptr(), total_sz, \
+      slice_off, slice_sz);
+
+  switch (ws) {
+    case 2:
+      LAUNCH(2);
+      break;
+    case 4:
+      LAUNCH(4);
+      break;
+    case 6:
+      LAUNCH(6);
+      break;
+    case 8:
+      LAUNCH(8);
+      break;
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef LAUNCH
+}
+
+// ---------------------------------------------------------------------------
+// Python binding
+// ---------------------------------------------------------------------------
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_reduce_scatter", &moe_reduce_scatter,
+        "Lamport MoE reduce-scatter");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}

csrc/standalone/custom_ar/moe_reduce_scatter.py

@@ -0,0 +1,102 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Lamport-based MoE reduce-scatter for EP combine.
+
+JIT-compiles the CUDA kernel on first use (cached afterwards).
+Uses the same Lamport sentinel protocol as the all-gather kernel.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+import torch
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_reduce_scatter.cu"))
+    _lib = load(
+        name="moe_reduce_scatter_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_RS_VERBOSE", "") == "1",
+    )
+    return _lib
+
+
+class MoeReduceScatter:
+    """Lamport-based MoE combine reduce-scatter with double buffering."""
+
+    def __init__(self, ca_comm):
+        self.rank = ca_comm.rank
+        self.world_size = ca_comm.world_size
+        self.device = ca_comm.device
+        self.buffer_ptrs = ca_comm.buffer_ptrs
+        self.max_size = ca_comm.max_size
+
+        ws = self.world_size
+        # Use the SECOND half of the IPC buffer (first half reserved for
+        # MoeAllGather) to avoid overlapping writes.
+        half_size = (self.max_size // 2) & ~15
+        self.buffer_offset = half_size
+
+        # Double-buffer layout within our half: 2 segments, each ws rank-slots.
+        self.seg_capacity = (half_size // 2) & ~15
+        self.rank_stride = (self.seg_capacity // ws) & ~15
+        self.max_per_rank = self.rank_stride
+
+        # Buffer pointer array on device — offset to our half.
+        self._buf_ptrs = torch.zeros(
+            8, dtype=torch.int64, device=f"cuda:{self.device.index}"
+        )
+        for i in range(ws):
+            self._buf_ptrs[i] = self.buffer_ptrs[i] + self.buffer_offset
+        self._buf_ptrs_ptr = self._buf_ptrs.data_ptr()
+
+        # Counters: [0]=unused, [1]=seg (0/1), [2]=prev_total_sz.
+        self._counters = torch.zeros(
+            3, dtype=torch.int32, device=f"cuda:{self.device.index}"
+        )
+        self._counters_ptr = self._counters.data_ptr()
+
+        # Initialize our half with sentinels.
+        lib = _load_lib()
+        lib.lamport_init(self.buffer_ptrs[self.rank] + self.buffer_offset, half_size)
+        torch.accelerator.synchronize(self.device)
+
+    def reduce_scatter(
+        self,
+        input: torch.Tensor,
+    ) -> torch.Tensor:
+        """Reduce-scatter input [N_total, D] bf16 → output [N_per_rank, D] bf16."""
+        lib = _load_lib()
+        ws = self.world_size
+
+        assert input.dim() == 2
+        N_total, D = input.shape
+        assert N_total % ws == 0
+        N_per_rank = N_total // ws
+
+        output = torch.empty((N_per_rank, D), dtype=input.dtype, device=input.device)
+
+        lib.moe_reduce_scatter(
+            self._buf_ptrs_ptr,
+            self._counters_ptr,
+            self.rank,
+            self.world_size,
+            self.seg_capacity,
+            self.rank_stride,
+            input,
+            output,
+        )
+        return output

csrc/standalone/custom_ar/moe_rs_fused.cu

@@ -0,0 +1,308 @@
+// Lamport reduce-scatter fused with residual add + RMSNorm.
+//
+// Replaces three separate kernels (RS + residual_add + RMSNorm) with one:
+//   1. PUSH: write MoE output to all peers' Lamport buffers.
+//   2. CLEAR: write sentinels to old segment.
+//   3. POLL+REDUCE+FUSE (per-token):
+//      a. Volatile-load from all peers, sum in fp32.
+//      b. Add residual.
+//      c. Compute RMSNorm (block reduction for variance).
+//      d. Store normed output + updated residual.
+//
+// Saves: one kernel launch (~3-5µs) + one global memory round-trip per layer.
+
+#include <cuda.h>
+#include <cuda_bf16.h>
+#include <cuda_runtime.h>
+#include <torch/extension.h>
+
+#include <c10/cuda/CUDAGuard.h>
+#include <c10/cuda/CUDAStream.h>
+
+#define DINLINE __device__ __forceinline__
+
+constexpr uint32_t SENTINEL = 0x80000000u;
+constexpr int kMaxBlocks = 36;
+
+// Each token has D=7168 bf16 values = 896 int4 vectors.
+// With 512 threads: ceil(896/512) = 2 int4 per thread = 16 fp32 values.
+constexpr int kMaxValsPerThread = 16;
+
+static DINLINE int4 ld128v(const void* addr) {
+  int4 v;
+  asm volatile("ld.volatile.global.v4.b32 {%0,%1,%2,%3}, [%4];"
+               : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w)
+               : "l"(addr));
+  return v;
+}
+
+static DINLINE bool has_sentinel(int4 v) {
+  return reinterpret_cast<uint32_t&>(v.x) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.y) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.z) == SENTINEL |
+         reinterpret_cast<uint32_t&>(v.w) == SENTINEL;
+}
+
+static DINLINE int4 remove_sentinel(int4 v) {
+  if (reinterpret_cast<uint32_t&>(v.x) == SENTINEL) v.x = 0;
+  if (reinterpret_cast<uint32_t&>(v.y) == SENTINEL) v.y = 0;
+  if (reinterpret_cast<uint32_t&>(v.z) == SENTINEL) v.z = 0;
+  if (reinterpret_cast<uint32_t&>(v.w) == SENTINEL) v.w = 0;
+  return v;
+}
+
+// bf16 helpers
+static DINLINE void accumulate_bf16(float* acc, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) acc[k] += __bfloat162float(bp[k]);
+}
+
+static DINLINE void add_bf16_to_fp32(float* dst, int4 v) {
+  const __nv_bfloat16* bp = reinterpret_cast<const __nv_bfloat16*>(&v);
+#pragma unroll
+  for (int k = 0; k < 8; k++) dst[k] += __bfloat162float(bp[k]);
+}
+
+static DINLINE int4 fp32_to_bf16_int4(const float* vals) {
+  int4 out;
+  __nv_bfloat16* bp = reinterpret_cast<__nv_bfloat16*>(&out);
+#pragma unroll
+  for (int k = 0; k < 8; k++) bp[k] = __float2bfloat16(vals[k]);
+  return out;
+}
+
+// Block-level tree reduction in shared memory.
+static DINLINE float block_reduce_sum(float val, float* smem) {
+  smem[threadIdx.x] = val;
+  __syncthreads();
+  for (int s = blockDim.x / 2; s > 0; s >>= 1) {
+    if (threadIdx.x < s) smem[threadIdx.x] += smem[threadIdx.x + s];
+    __syncthreads();
+  }
+  return smem[0];
+}
+
+// ---------------------------------------------------------------------------
+// Fused reduce-scatter + residual + RMSNorm kernel
+// ---------------------------------------------------------------------------
+
+template <int ngpus>
+__global__ void __launch_bounds__(512, 1) moe_rs_fused_kernel(
+    int64_t* buf_ptrs, int* counters, int rank, int seg_capacity,
+    int rank_stride,
+    const void* input,        // [N_total, D] bf16 — MoE output
+    const void* residual_in,  // [N_per_rank, D] bf16 — skip connection
+    const void* gamma,        // [D] bf16 — RMSNorm weight
+    void* normed_out,         // [N_per_rank, D] bf16 — normed result
+    void* residual_out,       // [N_per_rank, D] bf16 — updated residual
+    int total_sz,             // int4 units of full input
+    int slice_off,            // int4 offset to this rank's slice
+    int slice_sz,             // int4 units of this rank's slice
+    int D_int4,               // int4 units per token (hidden_dim * 2 / 16)
+    int N_per_rank,           // tokens in this rank's slice
+    float eps) {              // RMSNorm epsilon
+  using V = int4;
+  const int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  const int stride = gridDim.x * blockDim.x;
+
+  const int seg = counters[1];
+  const int prev_total_sz = counters[2];
+  const int cur_seg = seg;
+  const int old_seg = 1 - seg;
+
+  char* bufs[ngpus];
+#pragma unroll
+  for (int r = 0; r < ngpus; r++)
+    bufs[r] = reinterpret_cast<char*>(buf_ptrs[r]) + cur_seg * seg_capacity;
+
+  V sent;
+  sent.x = sent.y = sent.z = sent.w = static_cast<int>(SENTINEL);
+
+  // ---- Phase 1: PUSH full MoE output to ALL peers ----
+  {
+    const V* src = reinterpret_cast<const V*>(input);
+    for (int i = tid; i < total_sz; i += stride) {
+      V val = remove_sentinel(src[i]);
+#pragma unroll
+      for (int r = 0; r < ngpus; r++)
+        reinterpret_cast<V*>(bufs[r] + rank * rank_stride)[i] = val;
+    }
+  }
+
+  // ---- Phase 2: CLEAR old segment ----
+  if (prev_total_sz > 0) {
+    char* clr_base =
+        reinterpret_cast<char*>(buf_ptrs[rank]) + old_seg * seg_capacity;
+#pragma unroll
+    for (int r = 0; r < ngpus; r++) {
+      V* clr = reinterpret_cast<V*>(clr_base + r * rank_stride);
+      for (int i = tid; i < prev_total_sz; i += stride) clr[i] = sent;
+    }
+  }
+
+  // ---- Phase 3: Fused POLL + REDUCE + RESIDUAL + RMSNORM ----
+  // Each block handles one token. Only first N_per_rank blocks participate.
+  if (blockIdx.x < N_per_rank) {
+    const int token = blockIdx.x;
+    const int token_off = slice_off + token * D_int4;  // in the full buffer
+    char* my = bufs[rank];
+
+    extern __shared__ float smem[];
+
+    // Register storage for intermediate fp32 values.
+    float local_vals[kMaxValsPerThread];
+    int n_vals = 0;
+    float partial_sum_sq = 0.0f;
+
+    // Pass 1: poll + reduce + add residual + compute sum_sq
+    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
+      float acc[8] = {0, 0, 0, 0, 0, 0, 0, 0};
+
+      // Poll all ranks' data for this position.
+#pragma unroll
+      for (int s = 0; s < ngpus; s++) {
+        V val;
+        do {
+          val = ld128v(reinterpret_cast<V*>(my + s * rank_stride) + token_off +
+                       pos);
+        } while (has_sentinel(val));
+        accumulate_bf16(acc, val);
+      }
+
+      // Add residual.
+      V res = reinterpret_cast<const V*>(residual_in)[token * D_int4 + pos];
+      add_bf16_to_fp32(acc, res);
+
+      // Store in registers and accumulate sum_sq.
+#pragma unroll
+      for (int k = 0; k < 8; k++) {
+        local_vals[n_vals++] = acc[k];
+        partial_sum_sq += acc[k] * acc[k];
+      }
+    }
+
+    // Block-level reduction: total sum of squares.
+    float total_sum_sq = block_reduce_sum(partial_sum_sq, smem);
+    float rms_scale = rsqrtf(total_sum_sq / (D_int4 * 8) + eps);
+
+    // Pass 2: apply RMSNorm, store outputs.
+    n_vals = 0;
+    const V* gamma_v = reinterpret_cast<const V*>(gamma);
+    for (int pos = threadIdx.x; pos < D_int4; pos += blockDim.x) {
+      V gv = gamma_v[pos];
+      const __nv_bfloat16* gp = reinterpret_cast<const __nv_bfloat16*>(&gv);
+
+      // Build normed output and residual output.
+      float normed_fp32[8], res_fp32[8];
+#pragma unroll
+      for (int k = 0; k < 8; k++) {
+        float val = local_vals[n_vals++];
+        res_fp32[k] = val;  // residual_out
+        normed_fp32[k] = val * rms_scale * __bfloat162float(gp[k]);  // normed
+      }
+
+      reinterpret_cast<V*>(normed_out)[token * D_int4 + pos] =
+          fp32_to_bf16_int4(normed_fp32);
+      reinterpret_cast<V*>(residual_out)[token * D_int4 + pos] =
+          fp32_to_bf16_int4(res_fp32);
+    }
+  }
+
+  // ---- Phase 4: ADVANCE ----
+  if (blockIdx.x == 0 && threadIdx.x == 0) {
+    counters[1] = 1 - seg;
+    counters[2] = total_sz;
+  }
+}
+
+// ---------------------------------------------------------------------------
+// Sentinel init + host launcher
+// ---------------------------------------------------------------------------
+
+__global__ void lamport_init_kernel(uint32_t* buf, int n) {
+  int tid = blockIdx.x * blockDim.x + threadIdx.x;
+  int stride = gridDim.x * blockDim.x;
+  for (int i = tid; i < n; i += stride) buf[i] = SENTINEL;
+}
+
+void lamport_init(int64_t buf_ptr, int64_t nbytes) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  int n = static_cast<int>(nbytes / 4);
+  lamport_init_kernel<<<256, 256, 0, stream>>>(
+      reinterpret_cast<uint32_t*>(buf_ptr), n);
+}
+
+void moe_rs_fused(int64_t buf_ptrs_ptr, int64_t counters_ptr, int64_t rank,
+                  int64_t world_size, int64_t seg_capacity, int64_t rank_stride,
+                  torch::Tensor input,         // [N_total, D] bf16
+                  torch::Tensor residual_in,   // [N_per_rank, D] bf16
+                  torch::Tensor gamma,         // [D] bf16
+                  torch::Tensor normed_out,    // [N_per_rank, D] bf16
+                  torch::Tensor residual_out,  // [N_per_rank, D] bf16
+                  double eps) {
+  auto stream = c10::cuda::getCurrentCUDAStream().stream();
+  TORCH_CHECK(input.is_contiguous() && residual_in.is_contiguous());
+  TORCH_CHECK(gamma.is_contiguous() && normed_out.is_contiguous());
+  TORCH_CHECK(residual_out.is_contiguous());
+  TORCH_CHECK(input.scalar_type() == torch::kBFloat16);
+
+  int ws = static_cast<int>(world_size);
+  int r = static_cast<int>(rank);
+  int64_t N_total = input.size(0);
+  int64_t D = input.size(1);
+  TORCH_CHECK(N_total % ws == 0);
+  int N_per_rank = static_cast<int>(N_total / ws);
+
+  int64_t input_bytes = input.numel() * input.element_size();
+  TORCH_CHECK(input_bytes % 16 == 0);
+  TORCH_CHECK(input_bytes <= rank_stride);
+
+  int total_sz = static_cast<int>(input_bytes / 16);
+  int D_int4 = static_cast<int>(D * 2 / 16);  // bf16 elements → int4 units
+  int slice_sz = total_sz / ws;
+  int slice_off = r * slice_sz;
+
+  int threads = 512;
+  // Need at least N_per_rank blocks for Phase 3 (one per token).
+  int blocks = std::max(
+      N_per_rank, std::min(kMaxBlocks, (total_sz + threads - 1) / threads));
+
+  auto* bp = reinterpret_cast<int64_t*>(buf_ptrs_ptr);
+  auto* ct = reinterpret_cast<int*>(counters_ptr);
+  int sc = static_cast<int>(seg_capacity);
+  int rs = static_cast<int>(rank_stride);
+  int smem = threads * sizeof(float);
+
+#define LAUNCH(ng)                                                      \
+  moe_rs_fused_kernel<ng><<<blocks, threads, smem, stream>>>(           \
+      bp, ct, r, sc, rs, input.data_ptr(), residual_in.data_ptr(),      \
+      gamma.data_ptr(), normed_out.data_ptr(), residual_out.data_ptr(), \
+      total_sz, slice_off, slice_sz, D_int4, N_per_rank,                \
+      static_cast<float>(eps));
+
+  switch (ws) {
+    case 2:
+      LAUNCH(2);
+      break;
+    case 4:
+      LAUNCH(4);
+      break;
+    case 6:
+      LAUNCH(6);
+      break;
+    case 8:
+      LAUNCH(8);
+      break;
+    default:
+      TORCH_CHECK(false, "world_size must be 2, 4, 6, or 8");
+  }
+#undef LAUNCH
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("moe_rs_fused", &moe_rs_fused,
+        "Fused Lamport reduce-scatter + residual + RMSNorm");
+  m.def("lamport_init", &lamport_init,
+        "Initialize Lamport buffer with sentinels");
+}

csrc/standalone/custom_ar/moe_rs_fused.py

@@ -0,0 +1,29 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+JIT wrapper for the fused reduce-scatter + residual + RMSNorm kernel.
+"""
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+
+_lib = None
+
+
+def _load_lib():
+    global _lib
+    if _lib is not None:
+        return _lib
+
+    from torch.utils.cpp_extension import load
+
+    src = str(Path(__file__).with_name("moe_rs_fused.cu"))
+    _lib = load(
+        name="moe_rs_fused_kernel",
+        sources=[src],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=os.environ.get("MOE_RS_FUSED_VERBOSE", "") == "1",
+    )
+    return _lib

csrc/standalone/custom_ar/test_moe_allgather.py

@@ -0,0 +1,137 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test for the Lamport-based MoE all-gather kernel."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import MoeAllGather, _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather, _load_lib
+
+    _load_lib()
+    dist.barrier()
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    # Build a fake ca_comm-like object.
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    # meta_ptrs not needed for Lamport approach
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    errors = 0
+
+    # Test with various token counts.
+    for N in [1, 4, 16, 64]:
+        topk = 8
+        hd = 3584
+        sd = 448
+        ids = (
+            torch.arange(N * topk, dtype=torch.int32, device=dev) + rank * 1000
+        ).reshape(N, topk)
+        wt = torch.ones(N, topk, dtype=torch.float32, device=dev) * (rank + 1) * 0.1
+        hs = torch.full((N, hd), rank + 1, dtype=torch.uint8, device=dev)
+        sc = torch.full((N, sd), rank + 1, dtype=torch.uint8, device=dev)
+
+        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
+
+        for src in range(ws):
+            s, e = src * N, (src + 1) * N
+            exp_ids = (
+                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
+            ).reshape(N, topk)
+            if not torch.equal(ids_g[s:e], exp_ids):
+                print(f"[{rank}] FAIL ids src={src} N={N}")
+                errors += 1
+            exp_wt = torch.full(
+                (N, topk), (src + 1) * 0.1, dtype=torch.float32, device=dev
+            )
+            if not torch.allclose(wt_g[s:e], exp_wt):
+                print(f"[{rank}] FAIL wt src={src} N={N}")
+                errors += 1
+            exp_hs = torch.full((N, hd), src + 1, dtype=torch.uint8, device=dev)
+            if not torch.equal(hs_g[s:e], exp_hs):
+                print(f"[{rank}] FAIL hs src={src} N={N}")
+                errors += 1
+            exp_sc = torch.full((N, sd), src + 1, dtype=torch.uint8, device=dev)
+            if not torch.equal(sc_g[s:e], exp_sc):
+                print(f"[{rank}] FAIL sc src={src} N={N}")
+                errors += 1
+
+        # Without scales.
+        ids_g2, wt_g2, hs_g2, _ = ag.gather(ids, wt, hs)
+        for src in range(ws):
+            s, e = src * N, (src + 1) * N
+            exp_ids = (
+                torch.arange(N * topk, dtype=torch.int32, device=dev) + src * 1000
+            ).reshape(N, topk)
+            if not torch.equal(ids_g2[s:e], exp_ids):
+                print(f"[{rank}] FAIL no-sc ids src={src} N={N}")
+                errors += 1
+
+    dist.barrier()
+    print(
+        f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'} (ws={ws})"
+    )
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/standalone/custom_ar/test_moe_allgather_stress.py

@@ -0,0 +1,173 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Stress test: random data, check bitwise correctness against NCCL."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_allgather import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_allgather import MoeAllGather, _load_lib
+
+    _load_lib()
+    dist.barrier()
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    ag = MoeAllGather(ca)
+    dist.barrier()
+
+    topk = 8
+    hd = 3584
+    sd = 448
+    errors = 0
+    total_checks = 0
+    sentinel_collisions = 0
+
+    for trial in range(200):
+        # All ranks must use the same N for NCCL reference.
+        N_tensor = torch.randint(1, 65, (1,), device=dev)
+        dist.broadcast(N_tensor, src=0)
+        N = N_tensor.item()
+        # Random data including possible sentinel values
+        ids = torch.randint(0, 256, (N, topk), dtype=torch.int32, device=dev)
+        wt = torch.randn(N, topk, dtype=torch.float32, device=dev)
+        hs = torch.randint(0, 256, (N, hd), dtype=torch.uint8, device=dev)
+        sc = torch.randint(0, 256, (N, sd), dtype=torch.uint8, device=dev)
+
+        # Count sentinel patterns in hidden_states (as uint32 view)
+        hs_u32 = hs.view(torch.int32)
+        sentinel_collisions += (hs_u32 == 0x80000000).sum().item()
+
+        # Custom kernel
+        ids_g, wt_g, hs_g, sc_g = ag.gather(ids, wt, hs, sc)
+
+        # NCCL reference
+        ids_ref = torch.empty(N * ws, topk, dtype=torch.int32, device=dev)
+        wt_ref = torch.empty(N * ws, topk, dtype=torch.float32, device=dev)
+        hs_ref = torch.empty(N * ws, hd, dtype=torch.uint8, device=dev)
+        sc_ref = torch.empty(N * ws, sd, dtype=torch.uint8, device=dev)
+        dist.all_gather_into_tensor(ids_ref, ids)
+        dist.all_gather_into_tensor(wt_ref, wt)
+        dist.all_gather_into_tensor(hs_ref, hs)
+        dist.all_gather_into_tensor(sc_ref, sc)
+
+        # Compare
+        if not torch.equal(ids_g, ids_ref):
+            mismatches = (ids_g != ids_ref).sum().item()
+            if trial < 5 or mismatches > 0:
+                print(f"[{rank}] trial={trial} ids MISMATCH: {mismatches} elements")
+            errors += 1
+        if not torch.equal(wt_g, wt_ref):
+            # Check for -0 vs +0 differences
+            bit_diff = wt_g.view(torch.int32) != wt_ref.view(torch.int32)
+            neg_zero_mask = wt_ref.view(torch.int32) == 0x80000000
+            real_errors = bit_diff & ~neg_zero_mask
+            if real_errors.any():
+                print(
+                    f"[{rank}] trial={trial} wt MISMATCH (non-negzero): {real_errors.sum().item()}"
+                )
+                errors += 1
+        if not torch.equal(hs_g, hs_ref):
+            mismatches = (hs_g != hs_ref).sum().item()
+            # Check if mismatches are due to sentinel collision
+            hs_g_u32 = hs_g.view(torch.int32)
+            hs_ref_u32 = hs_ref.view(torch.int32)
+            diff_mask = hs_g_u32 != hs_ref_u32
+            sentinel_mask = (hs_ref_u32 == 0x80000000) & diff_mask
+            non_sentinel = diff_mask & ~sentinel_mask
+            if non_sentinel.any():
+                print(
+                    f"[{rank}] trial={trial} hs NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
+                )
+                errors += 1
+            elif sentinel_mask.any():
+                if trial < 3:
+                    print(
+                        f"[{rank}] trial={trial} hs sentinel collision: "
+                        f"{sentinel_mask.sum().item()} words (expected rare)"
+                    )
+        if not torch.equal(sc_g, sc_ref):
+            mismatches = (sc_g != sc_ref).sum().item()
+            sc_g_u32 = sc_g.view(torch.int32) if sc_g.numel() % 4 == 0 else None
+            if sc_g_u32 is not None:
+                sc_ref_u32 = sc_ref.view(torch.int32)
+                diff_mask = sc_g_u32 != sc_ref_u32
+                sentinel_mask = (sc_ref_u32 == 0x80000000) & diff_mask
+                non_sentinel = diff_mask & ~sentinel_mask
+                if non_sentinel.any():
+                    print(
+                        f"[{rank}] trial={trial} sc NON-SENTINEL MISMATCH: {non_sentinel.sum().item()}"
+                    )
+                    errors += 1
+
+        total_checks += 1
+
+    dist.barrier()
+    print(
+        f"[rank {rank}] {total_checks} trials, {errors} real errors, "
+        f"{sentinel_collisions} sentinel patterns in hs data. "
+        f"{'PASSED' if errors == 0 else 'FAILED'}"
+    )
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/standalone/custom_ar/test_moe_reduce_scatter.py

@@ -0,0 +1,204 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test + benchmark for Lamport MoE reduce-scatter."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+    if rank == 0:
+        from moe_reduce_scatter import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    if rank != 0:
+        from moe_reduce_scatter import _load_lib
+
+        _load_lib()
+    dist.barrier()
+    from moe_reduce_scatter import MoeReduceScatter
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+    rs = MoeReduceScatter(ca)
+    dist.barrier()
+
+    D = 7168  # DeepSeek V3 hidden_dim
+    errors = 0
+
+    # ---- Correctness tests ----
+    for N_per_rank in [1, 4, 16]:
+        N_total = N_per_rank * ws
+
+        # Each rank gets a deterministic input.
+        torch.manual_seed(42)
+        # All ranks create the SAME "ground truth" inputs for each rank.
+        all_inputs = [
+            torch.randn(N_total, D, dtype=torch.bfloat16, device=dev) for _ in range(ws)
+        ]
+        # This rank's input is all_inputs[rank].
+        my_input = all_inputs[rank]
+
+        # Custom reduce-scatter.
+        custom_out = rs.reduce_scatter(my_input)
+
+        # NCCL reference: reduce_scatter_tensor.
+        nccl_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        dist.reduce_scatter_tensor(nccl_out, my_input)
+
+        # bf16 summation order differs between our kernel and NCCL,
+        # giving ~1-2 ULP differences. Use generous tolerance.
+        max_diff = (custom_out.float() - nccl_out.float()).abs().max().item()
+        if not torch.allclose(custom_out, nccl_out, atol=0.125, rtol=0.01):
+            mismatches = (
+                ((custom_out.float() - nccl_out.float()).abs() > 0.125).sum().item()
+            )
+            print(
+                f"[{rank}] N_per_rank={N_per_rank} MISMATCH: "
+                f"max_diff={max_diff:.6f}, mismatches={mismatches}"
+            )
+            errors += 1
+        else:
+            if rank == 0:
+                print(f"  N_per_rank={N_per_rank}: PASS (max_diff={max_diff:.6f})")
+
+    # ---- Benchmark ----
+    if rank == 0:
+        print(f"\nworld_size={ws}, max_per_rank={rs.max_per_rank} bytes")
+        print(
+            f"{'config':<12} {'lamport':>10} {'lamp_graph':>10} "
+            f"{'nccl':>10} {'nccl_graph':>10} {'speedup':>8}"
+        )
+        print("-" * 65)
+
+    configs = [
+        ("1tok", 1),
+        ("2tok", 2),
+        ("4tok", 4),
+        ("8tok", 8),
+        ("16tok", 16),
+        ("32tok", 32),
+        ("64tok", 64),
+        ("128tok", 128),
+        ("256tok", 256),
+    ]
+
+    for name, N_per_rank in configs:
+        N_total = N_per_rank * ws
+        input_bytes = N_total * D * 2  # bf16
+        if input_bytes > rs.max_per_rank:
+            if rank == 0:
+                print(f"{name:<12} {'skip (too large)':>40}")
+            continue
+
+        inp = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        c_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        n_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+
+        def run_lamport():
+            rs.reduce_scatter(inp)
+
+        def run_nccl():
+            dist.reduce_scatter_tensor(n_out, inp)
+
+        from bench_moe_allgather import gpu_timer
+
+        lam_us = gpu_timer(run_lamport)
+        try:
+            lam_g_us = gpu_timer_graph(run_lamport)
+        except Exception:
+            lam_g_us = float("nan")
+
+        nccl_us = gpu_timer(run_nccl)
+        try:
+            nccl_g_us = gpu_timer_graph(run_nccl)
+        except Exception:
+            nccl_g_us = float("nan")
+
+        if rank == 0:
+            speedup = nccl_g_us / lam_g_us if lam_g_us > 0 else float("nan")
+            print(
+                f"{name:<12} {lam_us:>9.1f}µ {lam_g_us:>9.1f}µ "
+                f"{nccl_us:>9.1f}µ {nccl_g_us:>9.1f}µ {speedup:>7.2f}x"
+            )
+
+    dist.barrier()
+    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

csrc/standalone/custom_ar/test_rs_fused.py

@@ -0,0 +1,270 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Test + benchmark: fused RS + residual + RMSNorm vs separate kernels."""
+
+import ctypes
+import os
+import sys
+
+import torch
+import torch.distributed as dist
+
+_cudart = ctypes.CDLL("libcudart.so")
+IPC = 64
+
+
+def _cc(r):
+    if r:
+        raise RuntimeError(f"CUDA err {r}")
+
+
+def ipc_buf(sz, rank, ws):
+    p = ctypes.c_void_p()
+    _cc(_cudart.cudaMalloc(ctypes.byref(p), sz))
+    _cc(_cudart.cudaMemset(p, 0, sz))
+    _cc(_cudart.cudaDeviceSynchronize())
+    h = (ctypes.c_byte * IPC)()
+    _cc(_cudart.cudaIpcGetMemHandle(ctypes.byref(h), p))
+    ah = [None] * ws
+    dist.all_gather_object(ah, bytes(h))
+    ptrs = []
+    for i in range(ws):
+        if i == rank:
+            ptrs.append(p.value)
+        else:
+            hh = (ctypes.c_byte * IPC)(*ah[i])
+            pp = ctypes.c_void_p()
+            _cc(_cudart.cudaIpcOpenMemHandle(ctypes.byref(pp), hh, ctypes.c_uint(1)))
+            ptrs.append(pp.value)
+    return ptrs
+
+
+def rms_norm_ref(x, gamma, eps):
+    """Reference RMSNorm in fp32."""
+    xf = x.float()
+    rms = torch.rsqrt(xf.pow(2).mean(-1, keepdim=True) + eps)
+    return (xf * rms * gamma.float()).to(x.dtype)
+
+
+def gpu_timer(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        fn()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def gpu_timer_graph(fn, warmup=20, repeats=200):
+    for _ in range(warmup):
+        fn()
+    torch.cuda.synchronize()
+    g = torch.cuda.CUDAGraph()
+    with torch.cuda.graph(g):
+        fn()
+    for _ in range(5):
+        g.replay()
+    torch.cuda.synchronize()
+    s = torch.cuda.Event(enable_timing=True)
+    e = torch.cuda.Event(enable_timing=True)
+    s.record()
+    for _ in range(repeats):
+        g.replay()
+    e.record()
+    torch.cuda.synchronize()
+    return s.elapsed_time(e) / repeats * 1000
+
+
+def main():
+    dist.init_process_group("nccl")
+    rank = dist.get_rank()
+    ws = dist.get_world_size()
+    torch.cuda.set_device(rank)
+    dev = f"cuda:{rank}"
+
+    sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+
+    # Compile fused kernel
+    if rank == 0:
+        from torch.utils.cpp_extension import load
+
+        load(
+            name="moe_rs_fused_kernel",
+            sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
+            extra_cuda_cflags=["-O3", "--use_fast_math"],
+            verbose=False,
+        )
+    dist.barrier()
+    from torch.utils.cpp_extension import load
+
+    fused_lib = load(
+        name="moe_rs_fused_kernel",
+        sources=[os.path.join(os.path.dirname(__file__), "moe_rs_fused.cu")],
+        extra_cuda_cflags=["-O3", "--use_fast_math"],
+        verbose=False,
+    )
+
+    # Also compile separate RS kernel for comparison
+    from moe_reduce_scatter import MoeReduceScatter
+
+    max_size = 8 * 1024 * 1024
+    bp = ipc_buf(max_size, rank, ws)
+    dist.barrier()
+
+    class FakeCA:
+        pass
+
+    ca = FakeCA()
+    ca.rank = rank
+    ca.world_size = ws
+    ca.device = torch.device(dev)
+    ca.buffer_ptrs = bp
+    ca.max_size = max_size
+
+    rs_separate = MoeReduceScatter(ca)
+
+    # Fused kernel setup (uses same buffer layout as MoeReduceScatter)
+    half_size = (max_size // 2) & ~15
+    buf_offset = half_size
+    fused_seg_cap = (half_size // 2) & ~15
+    fused_rank_stride = (fused_seg_cap // ws) & ~15
+
+    fused_buf_ptrs = torch.zeros(8, dtype=torch.int64, device=dev)
+    for i in range(ws):
+        fused_buf_ptrs[i] = bp[i] + buf_offset
+    fused_counters = torch.zeros(3, dtype=torch.int32, device=dev)
+
+    # Init sentinels for fused kernel's buffer region
+    fused_lib.lamport_init(bp[rank] + buf_offset, half_size)
+    torch.cuda.synchronize()
+    dist.barrier()
+
+    D = 7168
+    eps = 1e-6
+    gamma = torch.randn(D, dtype=torch.bfloat16, device=dev).abs() + 0.5
+    errors = 0
+
+    # ---- Correctness ----
+    for N_per_rank in [1, 4]:
+        N_total = N_per_rank * ws
+        torch.manual_seed(42 + rank)
+        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+
+        # Reference: NCCL RS + add + norm
+        rs_ref = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        dist.reduce_scatter_tensor(rs_ref, moe_out)
+        ref_residual = residual + rs_ref
+        ref_normed = rms_norm_ref(ref_residual, gamma, eps)
+
+        # Fused kernel
+        normed_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        residual_out = torch.empty(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        fused_lib.moe_rs_fused(
+            fused_buf_ptrs.data_ptr(),
+            fused_counters.data_ptr(),
+            rank,
+            ws,
+            fused_seg_cap,
+            fused_rank_stride,
+            moe_out,
+            residual,
+            gamma,
+            normed_out,
+            residual_out,
+            eps,
+        )
+        torch.cuda.synchronize()
+
+        # Compare
+        max_diff_res = (residual_out.float() - ref_residual.float()).abs().max().item()
+        max_diff_norm = (normed_out.float() - ref_normed.float()).abs().max().item()
+
+        ok = max_diff_res < 0.125 and max_diff_norm < 0.125
+        if rank == 0:
+            print(
+                f"  N_per_rank={N_per_rank}: {'PASS' if ok else 'FAIL'} "
+                f"(res_diff={max_diff_res:.4f}, norm_diff={max_diff_norm:.4f})"
+            )
+        if not ok:
+            errors += 1
+
+    # ---- Benchmark ----
+    if rank == 0:
+        print(f"\nBenchmark: D={D}, world_size={ws}")
+        print(
+            f"{'config':<10} {'fused':>10} {'fused_g':>10} "
+            f"{'RS+norm':>10} {'RS+norm_g':>10} {'speedup':>8}"
+        )
+        print("-" * 58)
+
+    for N_per_rank in [1, 2, 4, 8]:
+        N_total = N_per_rank * ws
+        input_bytes = N_total * D * 2
+        if input_bytes > fused_rank_stride:
+            if rank == 0:
+                print(f"{N_per_rank}tok      skip (too large)")
+            continue
+
+        moe_out = torch.randn(N_total, D, dtype=torch.bfloat16, device=dev)
+        residual = torch.randn(N_per_rank, D, dtype=torch.bfloat16, device=dev)
+        normed_out = torch.empty_like(residual)
+        residual_out = torch.empty_like(residual)
+        rs_out = torch.empty_like(residual)
+
+        # Fused
+        def run_fused():
+            fused_lib.moe_rs_fused(
+                fused_buf_ptrs.data_ptr(),
+                fused_counters.data_ptr(),
+                rank,
+                ws,
+                fused_seg_cap,
+                fused_rank_stride,
+                moe_out,
+                residual,
+                gamma,
+                normed_out,
+                residual_out,
+                eps,
+            )
+
+        # Separate: RS + add + norm (our Lamport RS + triton-like ops)
+        def run_separate():
+            rs_separate.reduce_scatter(moe_out)
+            # Simulate add + RMSNorm (in practice this is a fused triton kernel)
+            tmp = residual + rs_out
+            torch.rsqrt(tmp.float().pow(2).mean(-1, keepdim=True) + eps)
+
+        fused_us = gpu_timer(run_fused)
+        try:
+            fused_g = gpu_timer_graph(run_fused)
+        except Exception:
+            fused_g = float("nan")
+
+        sep_us = gpu_timer(run_separate)
+        try:
+            sep_g = gpu_timer_graph(run_separate)
+        except Exception:
+            sep_g = float("nan")
+
+        if rank == 0:
+            speedup = sep_g / fused_g if fused_g > 0 else float("nan")
+            print(
+                f"{N_per_rank}tok     {fused_us:>9.1f}µ {fused_g:>9.1f}µ "
+                f"{sep_us:>9.1f}µ {sep_g:>9.1f}µ {speedup:>7.2f}x"
+            )
+
+    dist.barrier()
+    print(f"[rank {rank}] {'PASSED' if errors == 0 else f'FAILED ({errors})'}")
+    dist.destroy_process_group()
+    return errors
+
+
+if __name__ == "__main__":
+    sys.exit(main())

vllm/config/compilation.py

@@ -727,6 +727,8 @@ class CompilationConfig:
         "vllm::kda_attention",
         "vllm::sparse_attn_indexer",
         "vllm::rocm_aiter_sparse_attn_indexer",
+        # For specialized models
+        "vllm::monolithic_attn",
     ]
 
     def compute_hash(self) -> str:

vllm/distributed/device_communicators/all2all.py

@@ -11,14 +11,305 @@ from vllm.distributed import get_dp_group, get_ep_group
 from vllm.forward_context import get_forward_context
 from vllm.logger import init_logger
 from vllm.platforms import current_platform
+from vllm.triton_utils import tl, triton
 from vllm.utils.flashinfer import (
     has_flashinfer_nvlink_one_sided,
     has_flashinfer_nvlink_two_sided,
 )
 from vllm.utils.import_utils import has_deep_ep, has_mori
+from vllm.utils.torch_utils import direct_register_custom_op
 
 from .base_device_communicator import All2AllManagerBase, Cache
 
+logger = init_logger(__name__)
+
+# ---- MoE dispatch custom op (torch.compile-safe) -----------------------
+# Registered as torch.ops.vllm.moe_dispatch so torch.compile treats it as
+# opaque.  The real implementation decides custom kernel vs NCCL at runtime
+# based on tensor sizes.
+
+# Global reference to the MoeAllGather/MoeReduceScatter per group, set lazily.
+_moe_ag_instances: dict[str, Any] = {}
+_moe_rs_instances: dict[str, Any] = {}
+
+
+def _get_ca_comm(group_name: str):
+    """Get the CustomAllreduce communicator for a group, or None."""
+    from vllm.distributed.parallel_state import _groups
+
+    group_ref = _groups.get(group_name)
+    if group_ref is None:
+        return None
+    group = group_ref()
+    if group is None:
+        return None
+    dc = getattr(group, "device_communicator", None)
+    if dc is None:
+        return None
+    ca = getattr(dc, "ca_comm", None)
+    if ca is None or ca.disabled or not getattr(ca, "fully_connected", False):
+        return None
+    return ca
+
+
+def _get_or_create_moe_ag(group_name: str):
+    """Get or lazily create the MoeAllGather for this group."""
+    if group_name in _moe_ag_instances:
+        return _moe_ag_instances[group_name]
+
+    ca = _get_ca_comm(group_name)
+    if ca is None:
+        _moe_ag_instances[group_name] = None
+        return None
+
+    try:
+        from vllm.jit_kernels.moe_all_gather.moe_allgather import (
+            MoeAllGather,
+        )
+
+        ag = MoeAllGather(ca)
+        _moe_ag_instances[group_name] = ag
+        logger.info("MoE custom all-gather initialized for group %s", group_name)
+        return ag
+    except ImportError:
+        _moe_ag_instances[group_name] = None
+        return None
+
+
+def _get_or_create_moe_rs(group_name: str):
+    """Get or lazily create the MoeReduceScatter for this group."""
+    if group_name in _moe_rs_instances:
+        return _moe_rs_instances[group_name]
+
+    ca = _get_ca_comm(group_name)
+    if ca is None:
+        _moe_rs_instances[group_name] = None
+        return None
+
+    try:
+        from vllm.jit_kernels.moe_reduce_scatter.moe_reduce_scatter import (
+            MoeReduceScatter,
+        )
+
+        rs = MoeReduceScatter(ca)
+        _moe_rs_instances[group_name] = rs
+        logger.info("MoE custom reduce-scatter initialized for group %s", group_name)
+        return rs
+    except ImportError:
+        _moe_rs_instances[group_name] = None
+        return None
+
+
+@triton.jit
+def _mask_topk_ids_kernel(
+    topk_ids_ptr,
+    topk_ids_stride,
+    topk,
+    num_actual_tokens_ptr,
+    BLOCK_SIZE: tl.constexpr,
+):
+    token_idx = tl.program_id(0)
+    num_actual_tokens = tl.load(num_actual_tokens_ptr)
+    if token_idx < num_actual_tokens:
+        return
+
+    block = tl.arange(0, BLOCK_SIZE)
+    mask = block < topk
+    tl.store(topk_ids_ptr + token_idx * topk_ids_stride + block, -1, mask=mask)
+
+
+def mask_topk_ids(topk_ids: torch.Tensor, num_actual_tokens: torch.Tensor) -> None:
+    num_tokens, topk = topk_ids.shape
+    _mask_topk_ids_kernel[(num_tokens,)](
+        topk_ids,
+        topk_ids.stride(0),
+        topk,
+        num_actual_tokens,
+        BLOCK_SIZE=triton.next_power_of_2(topk),
+    )
+
+
+def moe_dispatch(
+    tensors: list[torch.Tensor],
+    group_name: str,
+) -> list[torch.Tensor]:
+    """All-gather a list of tensors across the EP/DP group.
+
+    Tries the custom one-shot P2P kernel first; falls back to NCCL
+    all_gatherv if the custom kernel is unavailable or tensors are too large.
+    """
+    from vllm.distributed.parallel_state import _groups
+
+    group_ref = _groups.get(group_name)
+    assert group_ref is not None
+    group = group_ref()
+    assert group is not None
+
+    # Try custom kernel for small, uniform-sized batches.
+    n = len(tensors)
+    if 2 <= n <= 4:
+        forward_context = get_forward_context()
+        num_actual_tokens = getattr(forward_context, "num_actual_tokens", None)
+        if num_actual_tokens is not None:
+            topk_ids = tensors[2]
+            mask_topk_ids(topk_ids, num_actual_tokens)
+
+        moe_ag = _get_or_create_moe_ag(group_name)
+        if moe_ag is not None:
+            ws = moe_ag.world_size
+            # Compute packed size.
+            total_bytes = 0
+            ok = True
+            for t in tensors:
+                if not t.is_contiguous():
+                    ok = False
+                    break
+                nbytes = t.numel() * t.element_size()
+                if nbytes % 16 != 0:
+                    ok = False
+                    break
+                total_bytes = (total_bytes + 15) & ~15
+                total_bytes += nbytes
+
+            max_bytes = moe_ag.max_per_rank
+            if ok and total_bytes <= max_bytes:
+                outputs = [
+                    torch.empty(
+                        (t.shape[0] * ws, *t.shape[1:]),
+                        dtype=t.dtype,
+                        device=t.device,
+                    )
+                    for t in tensors
+                ]
+
+                from vllm.jit_kernels.moe_all_gather.moe_allgather import (
+                    _load_lib,
+                )
+
+                lib = _load_lib()
+                lib.moe_all_gather(
+                    moe_ag._buf_ptrs_ptr,
+                    moe_ag._counters_ptr,
+                    moe_ag.rank,
+                    ws,
+                    moe_ag.seg_capacity,
+                    moe_ag.rank_stride,
+                    tensors,
+                    outputs,
+                )
+                return outputs
+
+    # Fallback: NCCL all_gatherv.
+    return group.device_communicator.all_gatherv(tensors, dim=0)
+
+
+def moe_dispatch_fake(
+    tensors: list[torch.Tensor],
+    group_name: str,
+) -> list[torch.Tensor]:
+    """Fake impl for torch.compile tracing — returns empty tensors of the
+    correct gathered shape."""
+    from vllm.distributed.parallel_state import _groups
+
+    group_ref = _groups.get(group_name)
+    assert group_ref is not None
+    group = group_ref()
+    ws = group.world_size if group is not None else 1
+    return [
+        torch.empty((t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device)
+        for t in tensors
+    ]
+
+
+direct_register_custom_op(
+    op_name="moe_dispatch",
+    op_func=moe_dispatch,
+    fake_impl=moe_dispatch_fake,
+    mutates_args=[],
+)
+
+
+def moe_combine(
+    hidden_states: torch.Tensor,
+    group_name: str,
+) -> torch.Tensor:
+    """Reduce-scatter hidden_states across the EP/DP group.
+
+    Tries the custom Lamport kernel first; falls back to NCCL.
+    """
+    from vllm.distributed.parallel_state import _groups
+
+    group_ref = _groups.get(group_name)
+    assert group_ref is not None
+    group = group_ref()
+    assert group is not None
+
+    ws = group.world_size
+    n_total = hidden_states.shape[0]
+
+    # Try custom kernel for small, uniform, bf16 inputs.
+    if (
+        n_total % ws == 0
+        and hidden_states.is_contiguous()
+        and hidden_states.dtype == torch.bfloat16
+    ):
+        moe_rs = _get_or_create_moe_rs(group_name)
+        if moe_rs is not None:
+            input_bytes = hidden_states.numel() * hidden_states.element_size()
+            if input_bytes % 16 == 0 and input_bytes <= moe_rs.max_per_rank:
+                from vllm.jit_kernels.moe_reduce_scatter.moe_reduce_scatter import (  # noqa: E501
+                    _load_lib,
+                )
+
+                lib = _load_lib()
+                n_per_rank = n_total // ws
+                output = torch.empty(
+                    (n_per_rank, *hidden_states.shape[1:]),
+                    dtype=hidden_states.dtype,
+                    device=hidden_states.device,
+                )
+                lib.moe_reduce_scatter(
+                    moe_rs._buf_ptrs_ptr,
+                    moe_rs._counters_ptr,
+                    moe_rs.rank,
+                    ws,
+                    moe_rs.seg_capacity,
+                    moe_rs.rank_stride,
+                    hidden_states,
+                    output,
+                )
+                return output
+
+    # Fallback: NCCL reduce_scatter.
+    return group.device_communicator.reduce_scatter(hidden_states, dim=0)
+
+
+def moe_combine_fake(
+    hidden_states: torch.Tensor,
+    group_name: str,
+) -> torch.Tensor:
+    """Fake impl for torch.compile tracing."""
+    from vllm.distributed.parallel_state import _groups
+
+    group_ref = _groups.get(group_name)
+    assert group_ref is not None
+    group = group_ref()
+    ws = group.world_size if group is not None else 1
+    n_per_rank = hidden_states.shape[0] // ws
+    return torch.empty(
+        (n_per_rank, *hidden_states.shape[1:]),
+        dtype=hidden_states.dtype,
+        device=hidden_states.device,
+    )
+
+
+direct_register_custom_op(
+    op_name="moe_combine",
+    op_func=moe_combine,
+    fake_impl=moe_combine_fake,
+    mutates_args=[],
+)
+
 if has_flashinfer_nvlink_two_sided():
     from flashinfer.comm import Mapping  # type: ignore[import-not-found]
     from flashinfer.comm.mnnvl import MnnvlConfig  # type: ignore[import-not-found]
@@ -155,6 +446,132 @@ class AgRsAll2AllManager(All2AllManagerBase):
 
     def __init__(self, cpu_group, tcp_store_group=None):
         super().__init__(cpu_group, tcp_store_group)
+        self._moe_ag = None  # lazy: MoeAllGather instance
+
+    def _get_moe_allgather(self, dist_group):
+        """Lazily build the standalone MoE all-gather kernel wrapper."""
+        if self._moe_ag is not None:
+            return self._moe_ag
+
+        import logging
+
+        log = logging.getLogger(__name__)
+
+        # Check prerequisites: intra-node, NVLink, ca_comm available.
+        if self.internode:
+            log.warning("MoE AG: skipped (internode)")
+            return None
+        dc = getattr(dist_group, "device_communicator", None)
+        if dc is None:
+            log.warning("MoE AG: skipped (no device_communicator)")
+            return None
+        ca = getattr(dc, "ca_comm", None)
+        if ca is None:
+            log.warning("MoE AG: skipped (ca_comm is None)")
+            return None
+        if ca.disabled:
+            log.warning("MoE AG: skipped (ca_comm disabled)")
+            return None
+        if not getattr(ca, "fully_connected", False):
+            log.warning("MoE AG: skipped (not fully_connected)")
+            return None
+
+        try:
+            from vllm.jit_kernels.moe_all_gather.moe_allgather import (  # type: ignore
+                MoeAllGather,
+            )
+        except ImportError as e:
+            log.warning("MoE AG: skipped (import failed: %s)", e)
+            return None
+
+        self._moe_ag = MoeAllGather(ca)
+        log.info("MoE AG: custom kernel initialized successfully")
+        return self._moe_ag
+
+    # ---- helpers --------------------------------------------------------
+
+    def _custom_all_gather(
+        self,
+        dist_group,
+        tensors: list[torch.Tensor],
+        sizes: list[int],
+    ) -> list[torch.Tensor] | None:
+        """Try the one-shot P2P MoE all-gather kernel.  Returns None on
+        fallback (kernel unavailable, too many tensors, etc.)."""
+        import logging
+
+        _log = logging.getLogger(__name__)
+        _dbg = getattr(self, "_ag_dbg_count", 0)
+
+        n = len(tensors)
+        if n < 2 or n > 4:
+            if _dbg < 10:
+                _log.warning(f"MoE AG: n={n} (need 2-4)")
+                self._ag_dbg_count = _dbg + 1
+            return None
+        if any(s != sizes[0] for s in sizes[1:]):
+            if _dbg < 10:
+                _log.warning(f"MoE AG: non-uniform sizes {sizes}")
+                self._ag_dbg_count = _dbg + 1
+            return None
+        moe_ag = self._get_moe_allgather(dist_group)
+        if moe_ag is None:
+            return None
+
+        ws = moe_ag.world_size
+
+        total_bytes = 0
+        for t in tensors:
+            if not t.is_contiguous():
+                if _dbg < 10:
+                    _log.warning(f"MoE AG: non-contiguous {t.shape} {t.stride()}")
+                    self._ag_dbg_count = _dbg + 1
+                return None
+            nbytes = t.numel() * t.element_size()
+            if nbytes % 16 != 0:
+                if _dbg < 10:
+                    _log.warning(
+                        f"MoE AG: unaligned {nbytes}B shape={t.shape} dtype={t.dtype}"
+                    )
+                    self._ag_dbg_count = _dbg + 1
+                return None
+            total_bytes = (total_bytes + 15) & ~15
+            total_bytes += nbytes
+
+        if total_bytes > moe_ag.max_per_rank:
+            if _dbg < 10:
+                _log.warning(
+                    f"MoE AG: too large {total_bytes} > {moe_ag.max_per_rank} "
+                    f"n={n} shapes={[t.shape for t in tensors]}"
+                )
+                self._ag_dbg_count = _dbg + 1
+            return None
+
+        # Allocate per-tensor outputs.
+        outputs = [
+            torch.empty((t.shape[0] * ws, *t.shape[1:]), dtype=t.dtype, device=t.device)
+            for t in tensors
+        ]
+
+        from vllm.jit_kernels.moe_all_gather.moe_allgather import (  # type: ignore
+            _load_lib,
+        )
+
+        lib = _load_lib()
+        lib.moe_all_gather(
+            moe_ag._buf_ptrs_ptr,
+            moe_ag._counters_ptr,
+            moe_ag.rank,
+            ws,
+            moe_ag.seg_capacity,
+            moe_ag.rank_stride,
+            tensors,
+            outputs,
+        )
+
+        return outputs
+
+    # ---- public API -----------------------------------------------------
 
     def dispatch_router_logits(
         self,
@@ -176,19 +593,22 @@ class AgRsAll2AllManager(All2AllManagerBase):
         dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
 
-        tensors_to_gather = [hidden_states, router_logits]
+        tensors = [hidden_states, router_logits]
         if extra_tensors is not None:
-            tensors_to_gather.extend(extra_tensors)
+            tensors.extend(extra_tensors)
 
-        gathered_tensors = dist_group.all_gatherv(
-            tensors_to_gather,
-            dim=0,
-            sizes=sizes,
-        )
+        # Use custom op so torch.compile treats this as opaque.
+        # Runtime decides custom kernel vs NCCL based on tensor sizes.
+        if all(s == sizes[0] for s in sizes):
+            gathered = torch.ops.vllm.moe_dispatch(
+                tensors, group_name=dist_group.unique_name
+            )
+        else:
+            gathered = dist_group.all_gatherv(tensors, dim=0, sizes=sizes)
 
         if extra_tensors is not None:
-            return (gathered_tensors[0], gathered_tensors[1], gathered_tensors[2:])
-        return gathered_tensors[0], gathered_tensors[1]
+            return (gathered[0], gathered[1], gathered[2:])
+        return gathered[0], gathered[1]
 
     def dispatch(
         self,
@@ -202,7 +622,7 @@ class AgRsAll2AllManager(All2AllManagerBase):
         | tuple[torch.Tensor, torch.Tensor, torch.Tensor, list[torch.Tensor]]
     ):
         """
-        Gather hidden_states and router_logits from all dp ranks.
+        Gather hidden_states, topk_weights, topk_ids from all dp ranks.
         """
         dp_metadata = get_forward_context().dp_metadata
         assert dp_metadata is not None
@@ -211,24 +631,24 @@ class AgRsAll2AllManager(All2AllManagerBase):
         dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
         assert sizes[dist_group.rank_in_group] == hidden_states.shape[0]
 
-        tensors_to_gather = [hidden_states, topk_weights, topk_ids]
+        tensors = [hidden_states, topk_weights, topk_ids]
         if extra_tensors is not None:
-            tensors_to_gather.extend(extra_tensors)
+            tensors.extend(extra_tensors)
 
-        gathered_tensors = dist_group.all_gatherv(
-            tensors_to_gather,
-            dim=0,
-            sizes=sizes,
-        )
+        if all(s == sizes[0] for s in sizes):
+            gathered = torch.ops.vllm.moe_dispatch(
+                tensors, group_name=dist_group.unique_name
+            )
+        else:
+            gathered = dist_group.all_gatherv(tensors, dim=0, sizes=sizes)
 
-        hidden_states = gathered_tensors[0]
-        topk_weights = gathered_tensors[1]
-        topk_ids = gathered_tensors[2]
+        hidden_states = gathered[0]
+        topk_weights = gathered[1]
+        topk_ids = gathered[2]
 
         if extra_tensors is None:
             return hidden_states, topk_weights, topk_ids
-
-        return hidden_states, topk_weights, topk_ids, gathered_tensors[3:]
+        return hidden_states, topk_weights, topk_ids, gathered[3:]
 
     def combine(
         self, hidden_states: torch.Tensor, is_sequence_parallel: bool = False
@@ -242,11 +662,21 @@ class AgRsAll2AllManager(All2AllManagerBase):
         assert sizes is not None
 
         dist_group = get_ep_group() if is_sequence_parallel else get_dp_group()
-        hidden_states = dist_group.reduce_scatterv(hidden_states, dim=0, sizes=sizes)
+
+        # Use custom op so torch.compile treats this as opaque.
+        if all(s == sizes[0] for s in sizes):
+            hidden_states = torch.ops.vllm.moe_combine(
+                hidden_states, group_name=dist_group.unique_name
+            )
+        else:
+            hidden_states = dist_group.reduce_scatterv(
+                hidden_states, dim=0, sizes=sizes
+            )
         return hidden_states
 
     def destroy(self):
-        pass
+        if self._moe_ag is not None:
+            self._moe_ag = None
 
 
 class DeepEPAll2AllManagerBase(All2AllManagerBase):

vllm/distributed/device_communicators/cuda_communicator.py

@@ -41,8 +41,7 @@ class CudaCommunicator(DeviceCommunicatorBase):
             global_ranks,
             global_world_size,
         )
-        if "tp" not in unique_name:
-            # custom allreduce or torch symm mem can be used only by tp
+        if not any(tag in unique_name for tag in ("tp", "dp", "ep")):
             use_custom_allreduce = False
             use_torch_symm_mem = False
             use_flashinfer_allreduce = False
@@ -50,8 +49,12 @@ class CudaCommunicator(DeviceCommunicatorBase):
             from vllm.distributed.parallel_state import _ENABLE_CUSTOM_ALL_REDUCE
 
             use_custom_allreduce = _ENABLE_CUSTOM_ALL_REDUCE
-            use_torch_symm_mem = envs.VLLM_ALLREDUCE_USE_SYMM_MEM
-            use_flashinfer_allreduce = envs.VLLM_ALLREDUCE_USE_FLASHINFER
+            use_torch_symm_mem = (
+                "tp" in unique_name and envs.VLLM_ALLREDUCE_USE_SYMM_MEM
+            )
+            use_flashinfer_allreduce = (
+                "tp" in unique_name and envs.VLLM_ALLREDUCE_USE_FLASHINFER
+            )
 
         self.use_custom_allreduce = use_custom_allreduce
         self.use_torch_symm_mem = use_torch_symm_mem
@@ -238,8 +241,6 @@ class CudaCommunicator(DeviceCommunicatorBase):
 
     def reduce_scatter(self, input_: torch.Tensor, dim: int = -1):
         world_size = self.world_size
-        pynccl_comm = self.pynccl_comm
-        assert pynccl_comm is not None
         if dim < 0:
             # Convert negative dim to positive.
             dim += input_.dim()
@@ -256,6 +257,8 @@ class CudaCommunicator(DeviceCommunicatorBase):
             output_shape, dtype=input_tensor.dtype, device=input_tensor.device
         )
 
+        pynccl_comm = self.pynccl_comm
+        assert pynccl_comm is not None
         pynccl_comm.reduce_scatter(output, input_tensor)
 
         # Reshape before returning

vllm/envs.py

@@ -218,6 +218,7 @@ if TYPE_CHECKING:
     VLLM_USE_FLASHINFER_MOE_MXFP4_MXFP8_CUTLASS: bool = False
     VLLM_ALLREDUCE_USE_SYMM_MEM: bool = True
     VLLM_ALLREDUCE_USE_FLASHINFER: bool = False
+    VLLM_USE_SPECIALIZED_MODELS: bool = False
     VLLM_TUNED_CONFIG_FOLDER: str | None = None
     VLLM_GPT_OSS_SYSTEM_TOOL_MCP_LABELS: set[str] = set()
     VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT: bool = False
@@ -1531,6 +1532,10 @@ environment_variables: dict[str, Callable[[], Any]] = {
     "VLLM_ALLREDUCE_USE_FLASHINFER": lambda: bool(
         int(os.getenv("VLLM_ALLREDUCE_USE_FLASHINFER", "0"))
     ),
+    # Whether to enable specialized model implementations when available.
+    "VLLM_USE_SPECIALIZED_MODELS": lambda: bool(
+        int(os.getenv("VLLM_USE_SPECIALIZED_MODELS", "0"))
+    ),
     # Experimental: use this to enable MCP tool calling for non harmony models
     "VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT": lambda: bool(
         int(os.getenv("VLLM_USE_EXPERIMENTAL_PARSER_CONTEXT", "0"))

vllm/forward_context.py

@@ -310,6 +310,7 @@ def override_forward_context(forward_context: ForwardContext | None):
 def set_forward_context(
     attn_metadata: Any,
     vllm_config: VllmConfig,
+    num_actual_tokens: torch.Tensor | None = None,
     num_tokens: int | None = None,
     num_tokens_across_dp: torch.Tensor | None = None,
     cudagraph_runtime_mode: CUDAGraphMode = CUDAGraphMode.NONE,
@@ -377,6 +378,8 @@ def set_forward_context(
         additional_kwargs,
         skip_compiled,
     )
+    if num_actual_tokens is not None:
+        forward_context.num_actual_tokens = num_actual_tokens
 
     try:
         with (

vllm/jit_kernels/moe_all_gather/moe_allgather.cu

@@ -0,0 +1 @@
+/home/woosuk/workspace/vllm/csrc/standalone/custom_ar/moe_allgather.cu

vllm/jit_kernels/moe_all_gather/moe_allgather.py

@@ -0,0 +1 @@
+/home/woosuk/workspace/vllm/csrc/standalone/custom_ar/moe_allgather.py

vllm/jit_kernels/moe_reduce_scatter/moe_reduce_scatter.cu

@@ -0,0 +1 @@
+/home/woosuk/workspace/vllm/csrc/standalone/custom_ar/moe_reduce_scatter.cu

vllm/jit_kernels/moe_reduce_scatter/moe_reduce_scatter.py

@@ -0,0 +1 @@
+/home/woosuk/workspace/vllm/csrc/standalone/custom_ar/moe_reduce_scatter.py

vllm/jit_kernels/moe_reduce_scatter/moe_rs_fused.cu

@@ -0,0 +1 @@
+/home/woosuk/workspace/vllm/csrc/standalone/custom_ar/moe_rs_fused.cu

vllm/jit_kernels/moe_reduce_scatter/moe_rs_fused.py

@@ -0,0 +1 @@
+/home/woosuk/workspace/vllm/csrc/standalone/custom_ar/moe_rs_fused.py

vllm/model_executor/model_loader/dummy_loader.py

@@ -43,6 +43,13 @@ class DummyModelLoader(BaseModelLoader):
                 # random values to the weights.
                 initialize_dummy_weights(layer, model_config)
 
+                # Some models build derived weights from loaded parameters instead of
+                # storing them in checkpoints. Rebuild those tensors for dummy load.
+                for layer in model.modules():
+                    fuse_indexer_weights = getattr(layer, "fuse_indexer_weights", None)
+                    if callable(fuse_indexer_weights):
+                        fuse_indexer_weights()
+
     def _process_online_quant_layer(
         self,
         layer: nn.Module,

vllm/model_executor/models/registry.py

@@ -1298,6 +1298,15 @@ ModelRegistry = _ModelRegistry(
     }
 )
 
+if envs.VLLM_USE_SPECIALIZED_MODELS:
+    from vllm.model_executor.specialized_models import get_specialized_models
+
+    for _arch, (_mod, _cls) in get_specialized_models().items():
+        ModelRegistry.models[_arch] = _LazyRegisteredModel(
+            module_name=_mod,
+            class_name=_cls,
+        )
+
 _T = TypeVar("_T")
 
 

vllm/model_executor/specialized_models/README.md

@@ -0,0 +1,19 @@
+# [Experimental] Specialized Models
+
+This directory contains experimental, hand-tuned implementations for a small number of selected models. Each subdirectory targets a specific combination of model architecture (including all tensor shapes), quantization scheme, attention backend, and hardware.
+
+For example, `deepseek_v3_2_nvfp4/` targets `nvidia/DeepSeek-V3.2-NVFP4` with FP8 FlashInfer sparse MLA on Blackwell GPUs.
+
+**To opt in, set `VLLM_USE_SPECIALIZED_MODELS=1`.** When enabled, vLLM will prefer a specialized implementation over the generic one if a match is available.
+
+## Development Philosophy
+
+These implementations prioritize iteration speed and checkpoint-specific performance over broad reuse. They may target a very narrow use case and are not expected to cover the full vLLM feature surface. Known limitations include:
+
+- Parallelism strategy support may be incomplete (e.g. TP only, no EP, or vice versa).
+- `torch.compile` compatibility may be limited or untested.
+- Behavior with checkpoint formats outside the intended target is unsupported.
+
+Also, code duplication across implementations is intentional — each model should be free to evolve and be optimized independently without risk of regressing another.
+
+Code here is experimental and may be short-lived. Generic features and anything intended for long-term support should live in `../models/`.

vllm/model_executor/specialized_models/__init__.py

@@ -0,0 +1,36 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+Specialized model implementations.
+
+Each entry maps a vLLM architecture name to a (module_path, class_name)
+tuple, exactly like the main model registry.  When
+``VLLM_USE_SPECIALIZED_MODELS=1`` the main registry merges these entries
+so they take priority over the generic implementations.
+
+To add a new specialized model:
+  1. Create a sub-package under this directory.
+  2. Add the architecture -> (module, class) mapping to ``_MODELS`` below.
+"""
+
+from __future__ import annotations
+
+# ── Model list ───────────────────────────────────────────────────────
+# Maps architecture name -> (fully-qualified module, class name).
+# When the flag is enabled, these override the corresponding entries
+# in the main registry.
+_MODELS: dict[str, tuple[str, str]] = {
+    "DeepseekV32ForCausalLM": (
+        "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4",
+        "DeepseekV32ForCausalLM",
+    ),
+    "DeepSeekMTPModel": (
+        "vllm.model_executor.specialized_models.deepseek_v3_2_nvfp4",
+        "DeepSeekMTP",
+    ),
+}
+
+
+def get_specialized_models() -> dict[str, tuple[str, str]]:
+    """Return the specialized model registry."""
+    return _MODELS

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/README.md

@@ -0,0 +1,34 @@
+# nvidia/DeepSeek-V3.2-NVFP4
+
+An optimized implementation for `nvidia/DeepSeek-V3.2-NVFP4` with FP8 FlashInfer MLA on Blackwell GPUs.
+
+The main win comes from aggressively fusing ops in the attention path, across the MLA and sparse-indexer boundary, which is critical for low latency.
+On top of manual fusions, the implementation uses `torch.compile` with vLLM's custom fusion passes to fuse remaining miscellaneous ops.
+It is compatible with piecewise CUDA graphs for prefill and full CUDA graphs for decode.
+
+TP and EP are supported; PP is not.
+MTP is supported.
+
+## Usage
+
+```bash
+export VLLM_USE_SPECIALIZED_MODELS=1
+export VLLM_USE_V2_MODEL_RUNNER=1
+export TRTLLM_ENABLE_PDL=1
+
+NUM_GPUS=4
+
+# With TP
+vllm serve nvidia/DeepSeek-V3.2-NVFP4 \
+    -tp 4 \
+    --compilation-config '{"max_cudagraph_capture_size": 1024}' \
+    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
+    --kernel-config.enable_flashinfer_autotune=False
+
+# With attention DP + MoE EP
+vllm serve nvidia/DeepSeek-V3.2-NVFP4 \
+    -dp $NUM_GPUS -ep \
+    --compilation-config '{"max_cudagraph_capture_size": 1024}' \
+    --speculative-config '{"method": "mtp", "num_speculative_tokens": 1}' \
+    --kernel-config.enable_flashinfer_autotune=False
+```

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/__init__.py

@@ -0,0 +1,8 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V3.2 model optimized for SM100 (Blackwell)."""
+
+from .model import DeepseekV32ForCausalLM
+from .mtp import DeepSeekMTP
+
+__all__ = ["DeepseekV32ForCausalLM", "DeepSeekMTP"]

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/kernels.py

@@ -0,0 +1,715 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+import torch
+
+from vllm.triton_utils import tl, triton
+
+
+@triton.jit
+def _rms_norm(x, w, eps, HIDDEN_SIZE: tl.constexpr):
+    x = x.to(tl.float32)
+    mean_sq = tl.sum(x * x, axis=0) / HIDDEN_SIZE
+    rrms = tl.rsqrt(mean_sq + eps)
+    w = w.to(tl.float32)
+    return (x * rrms) * w
+
+
+@triton.jit
+def _layer_norm(x, w, b, eps, mask, HIDDEN_SIZE: tl.constexpr):
+    x = x.to(tl.float32)
+    mean = tl.sum(x, axis=0) / HIDDEN_SIZE
+    diff = tl.where(mask, x - mean, 0.0)
+    var = tl.sum(diff * diff, axis=0) / HIDDEN_SIZE
+    rstd = tl.rsqrt(var + eps)
+
+    w = w.to(tl.float32)
+    b = b.to(tl.float32)
+    return (x - mean) * rstd * w + b
+
+
+@triton.jit
+def _rope(
+    base_ptr,
+    head_stride,
+    cos,
+    sin,
+    NUM_HEADS: tl.constexpr,
+    HALF_ROT_DIM: tl.constexpr,
+    START_OFFSET: tl.constexpr,
+    INTERLEAVED: tl.constexpr,
+):
+    head_offset = tl.arange(0, NUM_HEADS)
+    dim_offset = tl.arange(0, HALF_ROT_DIM)
+    base_ptr = base_ptr + head_offset[:, None] * head_stride + START_OFFSET
+    if INTERLEAVED:
+        x1 = tl.load(base_ptr + dim_offset * 2).to(tl.float32)
+        x2 = tl.load(base_ptr + dim_offset * 2 + 1).to(tl.float32)
+        tl.store(base_ptr + dim_offset * 2, x1 * cos - x2 * sin)
+        tl.store(base_ptr + dim_offset * 2 + 1, x2 * cos + x1 * sin)
+    else:
+        x1 = tl.load(base_ptr + dim_offset).to(tl.float32)
+        x2 = tl.load(base_ptr + dim_offset + HALF_ROT_DIM).to(tl.float32)
+        tl.store(base_ptr + dim_offset, x1 * cos - x2 * sin)
+        tl.store(base_ptr + dim_offset + HALF_ROT_DIM, x2 * cos + x1 * sin)
+
+
+@triton.jit
+def _get_cos_sin(
+    cos_sin_cache_ptr,
+    cos_sin_cache_stride,
+    pos,
+    HALF_ROT_DIM: tl.constexpr,
+):
+    block = tl.arange(0, HALF_ROT_DIM)
+    cos = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block)
+    cos = cos.to(tl.float32)
+    sin = tl.load(cos_sin_cache_ptr + pos * cos_sin_cache_stride + block + HALF_ROT_DIM)
+    sin = sin.to(tl.float32)
+    return cos, sin
+
+
+@triton.jit
+def _fp8_ue8m0_quantize(vals):
+    """Quantize float32 values to FP8 E4M3 with a ue8m0 (power-of-2) scale.
+
+    Returns (fp8_vals, scale) so the caller can store them or reuse the scale.
+    """
+    vals = vals.to(tl.float32)
+    amax = tl.max(tl.abs(vals))
+    scale = tl.div_rn(tl.maximum(amax, 1e-4), 448.0)
+    scale = tl.math.exp2(tl.math.ceil(tl.math.log2(scale)))
+    fp8_vals = tl.div_rn(vals, scale).to(tl.float8e4nv)
+    return fp8_vals, scale
+
+
+@triton.jit
+def _fp8_quant_and_cache_write(
+    vals,
+    mask,
+    slot_idx,
+    kv_cache_ptr,
+    kv_cache_scale_ptr,
+    cache_block_size,
+    cache_stride,
+    offsets,
+    HEAD_DIM: tl.constexpr,
+):
+    k_fp8, scale = _fp8_ue8m0_quantize(vals)
+
+    block_idx = slot_idx // cache_block_size
+    block_offset = slot_idx % cache_block_size
+    block_start = block_idx * cache_block_size * cache_stride
+
+    tl.store(
+        kv_cache_ptr + block_start + block_offset * HEAD_DIM + offsets,
+        k_fp8,
+        mask=mask,
+    )
+    scale_byte_off = block_start + cache_block_size * HEAD_DIM + block_offset * 4
+    tl.store(kv_cache_scale_ptr + scale_byte_off // 4, scale)
+
+
+@triton.jit
+def _fused_norm_rope_kernel(
+    pos_ptr,
+    # Q RMS norm
+    q_c_ptr,
+    q_c_stride,
+    q_rms_norm_w_ptr,
+    q_rms_eps,
+    q_c_out_ptr,
+    q_c_out_stride,
+    Q_DIM: tl.constexpr,
+    Q_BLOCK_SIZE: tl.constexpr,
+    # KV RMS norm
+    kv_ptr,
+    kv_stride,
+    kv_rms_norm_w_ptr,
+    kv_rms_eps,
+    KV_DIM: tl.constexpr,
+    # KV RoPE
+    kpe_ptr,
+    kpe_stride,
+    kpe_rope_cos_sin_cache_ptr,
+    kpe_rope_cos_sin_cache_stride,
+    KPE_HALF_ROT_DIM: tl.constexpr,
+    # Index K layer norm
+    index_k_ptr,
+    index_k_stride,
+    index_k_layer_norm_w_ptr,
+    index_k_layer_norm_bias_ptr,
+    index_k_layer_norm_eps,
+    INDEX_K_DIM: tl.constexpr,
+    INDEX_K_BLOCK_SIZE: tl.constexpr,
+    # Index K RoPE
+    index_k_rope_cos_sin_cache_ptr,
+    index_k_rope_cos_sin_cache_stride,
+    INDEX_K_HALF_ROT_DIM: tl.constexpr,
+    # Index K fp32 scratch buffer for layernorm → RoPE handoff
+    index_k_normed_ptr,
+    # Cache params (shared by indexer K and MLA)
+    slot_mapping_ptr,
+    # Index K FP8 cache
+    indexer_cache_ptr,
+    indexer_cache_scale_ptr,
+    indexer_cache_block_size,
+    indexer_cache_stride,
+    # MLA KV cache (concat kv_c_normed + k_pe_roped, uses slot_mapping_ptr)
+    mla_cache_ptr,
+    mla_cache_block_stride,
+    mla_cache_entry_stride,
+    MLA_CACHE_FP8: tl.constexpr,
+    mla_cache_scale_ptr,
+    # Top k indices
+    topk_indices_ptr,
+    topk_indices_stride,
+    TOPK: tl.constexpr,
+    TOPK_BLOCK_SIZE: tl.constexpr,
+):
+    pid = tl.program_id(0)
+    tok_idx = tl.program_id(1)
+    if pid == 3:
+        # Fill top k indices buffer with -1
+        for i in range(0, TOPK, TOPK_BLOCK_SIZE):
+            offset = i + tl.arange(0, TOPK_BLOCK_SIZE)
+            mask = offset < TOPK
+            tl.store(
+                topk_indices_ptr + tok_idx * topk_indices_stride + offset,
+                -1,
+                mask=mask,
+            )
+        return
+
+    if slot_mapping_ptr is None:
+        # Memory profiling run.
+        return
+    slot_idx = tl.load(slot_mapping_ptr + tok_idx)
+    if slot_idx < 0:
+        # Padding
+        return
+
+    if pid == 2:
+        # Q RMS norm
+        q_block = tl.arange(0, Q_BLOCK_SIZE)
+        q_mask = q_block < Q_DIM
+        q_c = tl.load(q_c_ptr + tok_idx * q_c_stride + q_block, mask=q_mask, other=0.0)
+        q_c_rms_w = tl.load(q_rms_norm_w_ptr + q_block, mask=q_mask)
+        q_c = _rms_norm(q_c, q_c_rms_w, q_rms_eps, Q_DIM)
+        tl.store(q_c_out_ptr + tok_idx * q_c_out_stride + q_block, q_c, mask=q_mask)
+    elif pid == 1:
+        # KV RMS Norm + KV RoPE + MLA concat_and_cache.
+        # Merged so the normed kv_c and RoPE'd k_pe can be written
+        # to the MLA KV cache directly without a separate kernel.
+
+        # KV RMS Norm (result stays in registers for MLA cache write)
+        kv_block = tl.arange(0, KV_DIM)
+        kv_c = tl.load(kv_ptr + tok_idx * kv_stride + kv_block)
+        kv_c_rms_w = tl.load(kv_rms_norm_w_ptr + kv_block)
+        kv_c = _rms_norm(kv_c, kv_c_rms_w, kv_rms_eps, KV_DIM)
+
+        # KV RoPE (interleaved) on k_pe — in registers only.
+        # k_pe is not needed after the cache write (MLA decode reads
+        # from kv_cache), so we skip writing back to kpe_ptr.
+        pos = tl.load(pos_ptr + tok_idx)
+        cos, sin = _get_cos_sin(
+            kpe_rope_cos_sin_cache_ptr,
+            kpe_rope_cos_sin_cache_stride,
+            pos,
+            KPE_HALF_ROT_DIM,
+        )
+        dim_off = tl.arange(0, KPE_HALF_ROT_DIM)
+        kpe_base = kpe_ptr + tok_idx * kpe_stride
+        x1 = tl.load(kpe_base + dim_off * 2).to(tl.float32)
+        x2 = tl.load(kpe_base + dim_off * 2 + 1).to(tl.float32)
+        r1 = x1 * cos - x2 * sin
+        r2 = x2 * cos + x1 * sin
+
+        # MLA concat_and_cache: write [kv_c_normed, k_pe_roped] to cache.
+        if mla_cache_entry_stride == 0:
+            return
+
+        mla_block_size = mla_cache_block_stride // mla_cache_entry_stride
+        mla_block_idx = slot_idx // mla_block_size
+        mla_block_off = slot_idx % mla_block_size
+        dst = (
+            mla_cache_ptr
+            + mla_block_idx * mla_cache_block_stride
+            + mla_block_off * mla_cache_entry_stride
+        )
+        # kv_c_normed (KV_DIM elements)
+        if MLA_CACHE_FP8:
+            scale = tl.load(mla_cache_scale_ptr)
+            kv_c_fp8 = (kv_c.to(tl.float32) / scale).to(tl.float8e4nv)
+            tl.store(dst + kv_block, kv_c_fp8)
+        else:
+            tl.store(dst + kv_block, kv_c)
+        # k_pe_roped (from registers, interleaved layout)
+        if MLA_CACHE_FP8:
+            tl.store(dst + KV_DIM + dim_off * 2, (r1 / scale).to(tl.float8e4nv))
+            tl.store(dst + KV_DIM + dim_off * 2 + 1, (r2 / scale).to(tl.float8e4nv))
+        else:
+            tl.store(dst + KV_DIM + dim_off * 2, r1)
+            tl.store(dst + KV_DIM + dim_off * 2 + 1, r2)
+    elif pid == 0:
+        # Fused: Index K LayerNorm + RoPE + FP8 quant + cache write.
+        # Eliminates the separate indexer_k_quant_and_cache kernel launch.
+
+        # 1. LayerNorm → fp32 temp buffer
+        index_k_block = tl.arange(0, INDEX_K_BLOCK_SIZE)
+        index_k_mask = index_k_block < INDEX_K_DIM
+        index_k = tl.load(
+            index_k_ptr + tok_idx * index_k_stride + index_k_block,
+            mask=index_k_mask,
+            other=0.0,
+        )
+        index_k_w = tl.load(index_k_layer_norm_w_ptr + index_k_block, mask=index_k_mask)
+        index_k_b = tl.load(
+            index_k_layer_norm_bias_ptr + index_k_block, mask=index_k_mask
+        )
+        normed = _layer_norm(
+            index_k,
+            index_k_w,
+            index_k_b,
+            index_k_layer_norm_eps,
+            index_k_mask,
+            INDEX_K_DIM,
+        )
+        # Write to a fp32 scratch buffer so RoPE can read the two
+        # halves without Triton pointer-aliasing issues.
+        scratch = index_k_normed_ptr + tok_idx * INDEX_K_DIM
+        tl.store(scratch + index_k_block, normed, mask=index_k_mask)
+
+        # 2. RoPE (neox / non-interleaved) on the full vector.
+        pos = tl.load(pos_ptr + tok_idx)
+        cos_full = tl.load(
+            index_k_rope_cos_sin_cache_ptr
+            + pos * index_k_rope_cos_sin_cache_stride
+            + index_k_block % INDEX_K_HALF_ROT_DIM,
+            mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
+            other=1.0,
+        ).to(tl.float32)
+        sin_full = tl.load(
+            index_k_rope_cos_sin_cache_ptr
+            + pos * index_k_rope_cos_sin_cache_stride
+            + INDEX_K_HALF_ROT_DIM
+            + index_k_block % INDEX_K_HALF_ROT_DIM,
+            mask=index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
+            other=0.0,
+        ).to(tl.float32)
+        # XOR with HALF swaps the first/second half of the rotation
+        # region to get each element's partner.
+        partner_offs = tl.where(
+            index_k_block < 2 * INDEX_K_HALF_ROT_DIM,
+            index_k_block ^ INDEX_K_HALF_ROT_DIM,
+            index_k_block,
+        )
+        full = tl.load(scratch + index_k_block, mask=index_k_mask)
+        # Atomic read for the partner: tl.atomic_add(ptr, 0) returns the
+        # current value with guaranteed store visibility, avoiding the
+        # Triton compiler's aliasing issue with different offset expressions.
+        zeros = tl.zeros([INDEX_K_BLOCK_SIZE], dtype=tl.float32)
+        partner = tl.atomic_add(scratch + partner_offs, zeros, mask=index_k_mask)
+        sign = tl.where(index_k_block < INDEX_K_HALF_ROT_DIM, -1.0, 1.0)
+        roped = full * cos_full + sign * partner * sin_full
+        result = tl.where(index_k_block < 2 * INDEX_K_HALF_ROT_DIM, roped, full)
+
+        # 3. FP8 quantize + cache write from registers.
+        #    No need to write back to index_k_ptr — the only consumer
+        #    (sparse_attn_indexer) reads from the cache, not index_k.
+        _fp8_quant_and_cache_write(
+            result,
+            index_k_mask,
+            slot_idx,
+            indexer_cache_ptr,
+            indexer_cache_scale_ptr,
+            indexer_cache_block_size,
+            indexer_cache_stride,
+            index_k_block,
+            INDEX_K_DIM,
+        )
+
+
+def fused_norm_rope(
+    positions: torch.Tensor,
+    q_c: torch.Tensor,
+    q_rms_norm_w: torch.Tensor,
+    q_rms_eps: float,
+    kv_c: torch.Tensor,
+    kv_rms_norm_w: torch.Tensor,
+    kv_rms_eps: float,
+    k_pe: torch.Tensor,
+    k_rope_cos_sin_cache: torch.Tensor,
+    index_k: torch.Tensor,
+    index_k_layer_norm_w: torch.Tensor,
+    index_k_layer_norm_bias: torch.Tensor,
+    index_k_layer_norm_eps: float,
+    index_k_rope_cos_sin_cache: torch.Tensor,
+    topk_indices_buffer: torch.Tensor,
+    # Cache params for fused writes (single slot_mapping for both caches)
+    slot_mapping: torch.Tensor | None = None,
+    indexer_k_cache: torch.Tensor | None = None,
+    mla_kv_cache: torch.Tensor | None = None,
+    mla_kv_cache_dtype: str = "auto",
+    mla_k_scale: torch.Tensor | None = None,
+) -> torch.Tensor:
+    assert positions.ndim == 1
+    assert q_c.ndim == 2
+    assert kv_c.ndim == 2
+    assert k_pe.ndim == 2
+    assert index_k.ndim == 2
+    assert topk_indices_buffer.ndim == 2
+
+    num_tokens = positions.shape[0]
+    q_dim = q_c.shape[-1]
+    kv_dim = kv_c.shape[-1]
+    index_k_dim = index_k.shape[-1]
+    topk = topk_indices_buffer.shape[-1]
+    device = positions.device
+
+    # --- Indexer K cache setup ---
+    if indexer_k_cache is not None:
+        assert slot_mapping is not None
+        idx_cache_scale_view = indexer_k_cache.view(torch.uint8).view(torch.float32)
+        idx_cache_block_size = indexer_k_cache.shape[1]
+        idx_cache_stride = indexer_k_cache.shape[2]
+        if indexer_k_cache.dtype == torch.uint8:
+            indexer_k_cache = indexer_k_cache.view(torch.float8_e4m3fn)
+    else:
+        idx_cache_scale_view = torch.empty(0, dtype=torch.float32, device=device)
+        indexer_k_cache = torch.empty(0, dtype=torch.float8_e4m3fn, device=device)
+        slot_mapping = torch.full((num_tokens,), -1, dtype=torch.int64, device=device)
+        idx_cache_block_size = 1
+        idx_cache_stride = 1
+
+    # --- MLA KV cache setup ---
+    mla_cache_fp8 = mla_kv_cache_dtype != "auto"
+    if mla_kv_cache is not None:
+        mla_block_stride = mla_kv_cache.stride(0)
+        mla_entry_stride = mla_kv_cache.stride(1)
+        if mla_cache_fp8 and mla_kv_cache.dtype == torch.uint8:
+            mla_kv_cache = mla_kv_cache.view(torch.float8_e4m3fn)
+        if mla_k_scale is None:
+            mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
+    else:
+        # Dummy values — pid 2 will skip the MLA cache write because
+        # slot_mapping is all -1.
+        mla_kv_cache = torch.empty(0, dtype=torch.bfloat16, device=device)
+        mla_block_stride = 0
+        mla_entry_stride = 0
+        mla_k_scale = torch.ones(1, dtype=torch.float32, device=device)
+
+    # fp32 scratch buffer for layernorm output → RoPE handoff.
+    index_k_normed = torch.empty(
+        num_tokens, index_k_dim, dtype=torch.float32, device=device
+    )
+
+    q_c_out = torch.empty_like(q_c)
+    _fused_norm_rope_kernel[(4, num_tokens)](
+        positions,
+        # Q RMS norm
+        q_c,
+        q_c.stride(0),
+        q_rms_norm_w,
+        q_rms_eps,
+        q_c_out,
+        q_c_out.stride(0),
+        q_dim,
+        triton.next_power_of_2(q_dim),
+        # KV RMS norm
+        kv_c,
+        kv_c.stride(0),
+        kv_rms_norm_w,
+        kv_rms_eps,
+        kv_dim,
+        # KV RoPE
+        k_pe,
+        k_pe.stride(0),
+        k_rope_cos_sin_cache,
+        k_rope_cos_sin_cache.stride(0),
+        k_rope_cos_sin_cache.shape[-1] // 2,
+        # Index K layer norm + RoPE + FP8 quant
+        index_k,
+        index_k.stride(0),
+        index_k_layer_norm_w,
+        index_k_layer_norm_bias,
+        index_k_layer_norm_eps,
+        index_k_dim,
+        triton.next_power_of_2(index_k_dim),
+        index_k_rope_cos_sin_cache,
+        index_k_rope_cos_sin_cache.stride(0),
+        index_k_rope_cos_sin_cache.shape[-1] // 2,
+        index_k_normed,
+        # Cache params
+        slot_mapping,
+        indexer_k_cache,
+        idx_cache_scale_view,
+        idx_cache_block_size,
+        idx_cache_stride,
+        # MLA KV cache (uses same slot_mapping)
+        mla_kv_cache,
+        mla_block_stride,
+        mla_entry_stride,
+        mla_cache_fp8,
+        mla_k_scale,
+        # Top k indices buffer
+        topk_indices_buffer,
+        topk_indices_buffer.stride(0),
+        topk,
+        TOPK_BLOCK_SIZE=1024,
+    )
+    return q_c_out
+
+
+@triton.jit
+def _fused_q_kernel(
+    pos_ptr,
+    # MQA query PE: RoPE + FP8 pack into output tail
+    q_pe_ptr,
+    q_pe_stride0,
+    q_pe_stride1,
+    NUM_Q_HEADS: tl.constexpr,
+    q_pe_cos_sin_ptr,
+    q_pe_cos_sin_stride,
+    Q_PE_HALF_ROT_DIM: tl.constexpr,
+    # Index Q RoPE
+    index_q_ptr,
+    index_q_stride0,
+    index_q_stride1,
+    NUM_INDEX_Q_HEADS: tl.constexpr,
+    index_q_cos_sin_ptr,
+    index_q_cos_sin_stride,
+    INDEX_Q_HALF_ROT_DIM: tl.constexpr,
+    # Index Q Quantize
+    index_q_fp8_ptr,
+    index_q_fp8_stride0,
+    index_q_fp8_stride1,
+    INDEX_Q_HEAD_DIM: tl.constexpr,
+    # MQA query pack: quantize ql_nope and RoPE+quantize q_pe into mqa_q_fp8
+    ql_nope_ptr,
+    ql_nope_stride0,
+    ql_nope_stride1,
+    mqa_q_fp8_ptr,
+    mqa_q_fp8_stride0,
+    mqa_q_fp8_stride1,
+    q_scale_ptr,
+    QL_NOPE_DIM: tl.constexpr,
+    QL_NOPE_BLOCK: tl.constexpr,
+    # Index weights
+    index_weights_ptr,
+    index_weights_stride,
+    index_weights_softmax_scale,
+    index_weights_head_scale,
+    index_weights_out_ptr,
+    index_weights_out_stride,
+):
+    pid = tl.program_id(0)
+    tok_idx = tl.program_id(1)
+    head_idx = tl.program_id(2)
+
+    if pid == 2:
+        # ql_nope quantize + pack into the front of mqa_q_fp8.
+        if 2 * head_idx >= NUM_Q_HEADS:
+            return
+
+        scale = tl.load(q_scale_ptr)
+        for local_head in range(2):
+            q_head_idx = head_idx * 2 + local_head
+            if q_head_idx < NUM_Q_HEADS:
+                ql_nope_off = tl.arange(0, QL_NOPE_BLOCK)
+                ql_nope_mask = ql_nope_off < QL_NOPE_DIM
+                ql_nope = tl.load(
+                    ql_nope_ptr
+                    + tok_idx * ql_nope_stride0
+                    + q_head_idx * ql_nope_stride1
+                    + ql_nope_off,
+                    mask=ql_nope_mask,
+                ).to(tl.float32)
+                ql_nope_fp8 = (ql_nope / scale).to(tl.float8e4nv)
+                tl.store(
+                    mqa_q_fp8_ptr
+                    + tok_idx * mqa_q_fp8_stride0
+                    + q_head_idx * mqa_q_fp8_stride1
+                    + ql_nope_off,
+                    ql_nope_fp8,
+                    mask=ql_nope_mask,
+                )
+        return
+    elif pid == 0:
+        # q_pe RoPE + quantize + pack into the tail of mqa_q_fp8.
+        if 2 * head_idx >= NUM_Q_HEADS:
+            return
+
+        pos = tl.load(pos_ptr + tok_idx)
+        cos, sin = _get_cos_sin(
+            q_pe_cos_sin_ptr,
+            q_pe_cos_sin_stride,
+            pos,
+            Q_PE_HALF_ROT_DIM,
+        )
+
+        scale = tl.load(q_scale_ptr)
+        for local_head in range(2):
+            q_head_idx = head_idx * 2 + local_head
+            if q_head_idx < NUM_Q_HEADS:
+                rot_off = tl.arange(0, Q_PE_HALF_ROT_DIM)
+                x1 = tl.load(
+                    q_pe_ptr
+                    + tok_idx * q_pe_stride0
+                    + q_head_idx * q_pe_stride1
+                    + rot_off * 2,
+                ).to(tl.float32)
+                x2 = tl.load(
+                    q_pe_ptr
+                    + tok_idx * q_pe_stride0
+                    + q_head_idx * q_pe_stride1
+                    + rot_off * 2
+                    + 1
+                ).to(tl.float32)
+                r1 = x1 * cos - x2 * sin
+                r2 = x2 * cos + x1 * sin
+                tl.store(
+                    mqa_q_fp8_ptr
+                    + tok_idx * mqa_q_fp8_stride0
+                    + q_head_idx * mqa_q_fp8_stride1
+                    + QL_NOPE_DIM
+                    + rot_off * 2,
+                    (r1 / scale).to(tl.float8e4nv),
+                )
+                tl.store(
+                    mqa_q_fp8_ptr
+                    + tok_idx * mqa_q_fp8_stride0
+                    + q_head_idx * mqa_q_fp8_stride1
+                    + QL_NOPE_DIM
+                    + rot_off * 2
+                    + 1,
+                    (r2 / scale).to(tl.float8e4nv),
+                )
+        return
+    elif pid == 1:
+        # Index Q RoPE
+        if head_idx >= NUM_INDEX_Q_HEADS:
+            return
+
+        pos = tl.load(pos_ptr + tok_idx)
+        cos, sin = _get_cos_sin(
+            index_q_cos_sin_ptr,
+            index_q_cos_sin_stride,
+            pos,
+            INDEX_Q_HALF_ROT_DIM,
+        )
+        _rope(
+            index_q_ptr + tok_idx * index_q_stride0 + head_idx * index_q_stride1,
+            0,
+            cos,
+            sin,
+            1,
+            INDEX_Q_HALF_ROT_DIM,
+            0,
+            False,
+        )
+
+        # Index Q Quantize
+        index_q_block = tl.arange(0, INDEX_Q_HEAD_DIM)
+        index_q = tl.load(
+            index_q_ptr
+            + tok_idx * index_q_stride0
+            + head_idx * index_q_stride1
+            + index_q_block
+        )
+
+        index_q_fp8, index_q_scale = _fp8_ue8m0_quantize(index_q)
+        tl.store(
+            index_q_fp8_ptr
+            + tok_idx * index_q_fp8_stride0
+            + head_idx * index_q_fp8_stride1
+            + index_q_block,
+            index_q_fp8,
+        )
+
+        # Index weights update
+        index_weights = tl.load(
+            index_weights_ptr + tok_idx * index_weights_stride + head_idx
+        )
+        index_weights = index_weights.to(tl.float32)
+        index_weights *= index_q_scale
+        index_weights *= index_weights_softmax_scale
+        index_weights *= index_weights_head_scale
+        tl.store(
+            index_weights_out_ptr + tok_idx * index_weights_out_stride + head_idx,
+            index_weights,
+        )
+
+
+def fused_q(
+    positions: torch.Tensor,
+    q_pe: torch.Tensor,
+    q_pe_cos_sin_cache: torch.Tensor,
+    index_q: torch.Tensor,
+    index_q_cos_sin_cache: torch.Tensor,
+    ql_nope: torch.Tensor,
+    q_scale: torch.Tensor,
+    # Index weights
+    index_weights: torch.Tensor,
+    index_weights_softmax_scale: float,
+    index_weights_head_scale: float,
+) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
+    assert positions.ndim == 1
+    assert q_pe.ndim == 3
+    assert q_pe_cos_sin_cache.ndim == 2
+    assert index_q.ndim == 3
+    assert index_q_cos_sin_cache.ndim == 2
+
+    num_tokens = positions.shape[0]
+    num_q_heads = q_pe.shape[1]
+    num_index_q_heads = index_q.shape[1]
+    index_q_head_dim = index_q.shape[2]
+    assert ql_nope.ndim == 3
+    assert ql_nope.shape[:2] == q_pe.shape[:2]
+    mqa_q_fp8 = torch.empty(
+        q_pe.shape[0],
+        q_pe.shape[1],
+        ql_nope.shape[2] + q_pe.shape[2],
+        dtype=torch.float8_e4m3fn,
+        device=q_pe.device,
+    )
+
+    index_q_fp8 = torch.empty_like(index_q, dtype=torch.float8_e4m3fn)
+    index_weights_out = torch.empty_like(index_weights, dtype=torch.float32)
+    _fused_q_kernel[(3, num_tokens, num_index_q_heads)](
+        positions,
+        q_pe,
+        q_pe.stride(0),
+        q_pe.stride(1),
+        num_q_heads,
+        q_pe_cos_sin_cache,
+        q_pe_cos_sin_cache.stride(0),
+        q_pe_cos_sin_cache.shape[-1] // 2,
+        index_q,
+        index_q.stride(0),
+        index_q.stride(1),
+        num_index_q_heads,
+        index_q_cos_sin_cache,
+        index_q_cos_sin_cache.stride(0),
+        index_q_cos_sin_cache.shape[-1] // 2,
+        index_q_fp8,
+        index_q_fp8.stride(0),
+        index_q_fp8.stride(1),
+        index_q_head_dim,
+        ql_nope,
+        ql_nope.stride(0),
+        ql_nope.stride(1),
+        mqa_q_fp8,
+        mqa_q_fp8.stride(0),
+        mqa_q_fp8.stride(1),
+        q_scale,
+        ql_nope.shape[2],
+        triton.next_power_of_2(ql_nope.shape[2]),
+        index_weights,
+        index_weights.stride(0),
+        index_weights_softmax_scale,
+        index_weights_head_scale,
+        index_weights_out,
+        index_weights_out.stride(0),
+        num_warps=1,  # TODO: Tune this
+    )
+    return index_q_fp8, index_weights_out, mqa_q_fp8

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/layer.py

@@ -0,0 +1,573 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""
+MLA attention and decoder layer for DeepSeek V3.2 on SM100 (Blackwell).
+
+MLAAttention:
+  KV cache update -> W_UK_T absorption -> sparse attn kernel -> W_UV up-proj
+  MLAAttention kept only as a registration stub for KV cache / backend.
+
+DecoderLayer:
+  Single decoder layer: norm -> attn -> norm -> MoE/MLP.
+"""
+
+from __future__ import annotations
+
+import torch
+from torch import nn
+
+from vllm.config import CacheConfig, VllmConfig, get_current_vllm_config
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.forward_context import get_forward_context
+from vllm.model_executor.layers.attention.mla_attention import MLAAttention
+from vllm.model_executor.layers.layernorm import LayerNorm, RMSNorm
+from vllm.model_executor.layers.linear import (
+    ColumnParallelLinear,
+    ReplicatedLinear,
+    RowParallelLinear,
+)
+from vllm.model_executor.layers.quantization import QuantizationConfig
+from vllm.model_executor.layers.rotary_embedding import get_rope
+from vllm.model_executor.layers.sparse_attn_indexer import SparseAttnIndexer
+from vllm.model_executor.models.deepseek_v2 import (
+    DeepseekV32IndexerCache,
+    yarn_get_mscale,
+)
+from vllm.platforms import current_platform
+from vllm.utils.torch_utils import direct_register_custom_op
+from vllm.v1.attention.backends.mla.indexer import get_max_prefill_buffer_size
+
+from .kernels import fused_norm_rope, fused_q
+from .sparse_indexer import sparse_attn_indexer
+
+
+def dsa(
+    positions: torch.Tensor,
+    q_c: torch.Tensor,
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    index_k: torch.Tensor,
+    index_weights: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    layer = get_forward_context().no_compile_layers[layer_name]
+    attn = layer.attn
+    mla = attn.mla_attn
+
+    attn_metadata = get_forward_context().attn_metadata
+    if not isinstance(attn_metadata, dict):
+        output.zero_()
+        return output
+
+    mla_attn_metadata = attn_metadata.get(mla.layer_name)
+    if mla_attn_metadata is None:
+        output.zero_()
+        return output
+
+    num_actual_toks = mla_attn_metadata.num_actual_tokens  # type: ignore[attr-defined]
+    if num_actual_toks == 0:
+        output.zero_()
+        return output
+
+    # Step 2. fused norm + rope + cache writes
+    slot_mapping = None
+    indexer_k_cache = None
+    mla_kv_cache = None
+    mla_k_scale = None
+    idx_meta = attn_metadata.get(attn.indexer_k_cache.prefix)
+    if idx_meta is not None:
+        slot_mapping = idx_meta.slot_mapping  # type: ignore[attr-defined]
+        indexer_k_cache = attn.indexer_k_cache.kv_cache
+        mla_kv_cache = attn.mla_attn.kv_cache
+        mla_k_scale = attn.mla_attn._k_scale
+
+    q_c = fused_norm_rope(
+        positions,
+        q_c,
+        attn.q_a_layernorm_weight,
+        layer.rms_norm_eps,
+        kv_c,
+        attn.kv_a_layernorm_weight,
+        attn.rms_norm_eps,
+        k_pe,
+        attn.rotary_emb.cos_sin_cache,
+        index_k,
+        attn.indexer_k_norm.weight,
+        attn.indexer_k_norm.bias,
+        attn.rms_norm_eps,
+        attn.indexer_rope_emb.cos_sin_cache,
+        attn.topk_indices_buffer,
+        slot_mapping=slot_mapping,
+        indexer_k_cache=indexer_k_cache,
+        mla_kv_cache=mla_kv_cache,
+        mla_kv_cache_dtype=attn.mla_attn.kv_cache_dtype,
+        mla_k_scale=mla_k_scale,
+    )
+
+    # Step 3. q_c -> index_q, q
+    step3_out = torch.mm(q_c, layer._fused_step3_q_w.T)
+    index_q, q = step3_out.split(layer._q_split_sizes, dim=-1)
+    index_q = index_q.view(-1, attn.index_n_heads, attn.index_head_dim)
+    q = q.view(-1, attn.num_local_heads, attn.qk_head_dim)
+
+    # Step 4. Q RoPE + W_UK_T absorption + FP8 packing
+    q_nope, q_pe = q.split(
+        [mla.qk_nope_head_dim, mla.qk_rope_head_dim],
+        dim=-1,
+    )
+    q_nope = q_nope.transpose(0, 1)
+    ql_nope = torch.bmm(q_nope, mla.W_UK_T)
+    ql_nope = ql_nope.transpose(0, 1)
+
+    index_q_fp8, index_weights, mqa_q = fused_q(
+        positions,
+        q_pe,
+        attn.rotary_emb.cos_sin_cache,
+        index_q,
+        attn.indexer_rope_emb.cos_sin_cache,
+        ql_nope,
+        mla._q_scale,
+        index_weights,
+        attn.indexer_softmax_scale,
+        attn.index_n_heads**-0.5,
+    )
+
+    # Steps 5-6. Sparse indexer + MLA sparse decode attention
+    sparse_attn_indexer(
+        attn.indexer_k_cache.prefix,
+        attn.indexer_k_cache.kv_cache,
+        index_q_fp8,
+        index_weights,
+        attn.topk_tokens,
+        attn.index_head_dim,
+        layer.max_model_len,
+        layer.indexer_workspace_size,
+        attn.topk_indices_buffer,
+    )
+
+    mqa_q = mqa_q[:num_actual_toks]
+    kv_cache = mla.kv_cache
+    if mla.kv_cache_dtype.startswith("fp8") and mla.kv_cache_dtype != "fp8_ds_mla":
+        kv_cache = kv_cache.view(torch.float8_e4m3fn)
+    attn_out, _ = mla.impl.forward_mqa(mqa_q, kv_cache, mla_attn_metadata, mla)
+    x = attn_out.view(-1, mla.num_heads, mla.kv_lora_rank).transpose(0, 1)
+
+    out = output[:num_actual_toks].view(-1, mla.num_heads, mla.v_head_dim)
+    out = out.transpose(0, 1)
+    torch.bmm(x, mla.W_UV, out=out)
+    return output
+
+
+def dsa_fake(
+    positions: torch.Tensor,
+    q_c: torch.Tensor,
+    kv_c: torch.Tensor,
+    k_pe: torch.Tensor,
+    index_k: torch.Tensor,
+    index_weights: torch.Tensor,
+    output: torch.Tensor,
+    layer_name: str,
+) -> torch.Tensor:
+    del positions, q_c, kv_c, k_pe, index_k, index_weights, layer_name
+    return output
+
+
+direct_register_custom_op(
+    op_name="monolithic_attn",
+    op_func=dsa,
+    fake_impl=dsa_fake,
+    mutates_args=["output"],
+    dispatch_key=current_platform.dispatch_key,
+)
+
+
+class DeepseekV32DecoderLayer(nn.Module):
+    """
+    Single decoder layer: norm -> attn -> norm -> MoE/MLP.
+    Norms are raw weight + direct kernel call.
+    Gate inlined as raw weight, experts kept as FusedMoE for quantization.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config,
+        layer_idx: int,
+        topk_indices_buffer: torch.Tensor,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        compilation_config = get_current_vllm_config().compilation_config
+        if prefix in compilation_config.static_forward_context:
+            raise ValueError(f"Duplicate layer name: {prefix}")
+        compilation_config.static_forward_context[prefix] = self
+
+        self.layer_name = prefix
+        self.layer_idx = layer_idx
+        self.hidden_size = config.hidden_size
+        self.rms_norm_eps = config.rms_norm_eps
+        self.q_lora_rank = config.q_lora_rank
+        self.kv_lora_rank = config.kv_lora_rank
+        self.qk_rope_head_dim = config.qk_rope_head_dim
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        cache_config = vllm_config.cache_config
+        quant_config = vllm_config.quant_config
+        parallel_config = vllm_config.parallel_config
+        self.indexer_workspace_size = get_max_prefill_buffer_size(vllm_config)
+        self.max_model_len = vllm_config.model_config.max_model_len
+
+        # Use the regular vLLM RMSNorm modules so the compiler sees the
+        # canonical residual-add + RMSNorm pattern.
+        dtype = torch.get_default_dtype()
+        self.input_layernorm = RMSNorm(
+            hidden_size=config.hidden_size,
+            eps=config.rms_norm_eps,
+            dtype=dtype,
+        )
+        self.post_attention_layernorm = RMSNorm(
+            hidden_size=config.hidden_size,
+            eps=config.rms_norm_eps,
+            dtype=dtype,
+        )
+
+        # Fused QKV A-projection lives inside self_attn namespace
+        # for weight loading compatibility with original checkpoint paths
+        from vllm.model_executor.models.deepseek_v2 import (
+            DeepSeekV2FusedQkvAProjLinear,
+        )
+
+        self.self_attn = nn.Module()
+        self.self_attn.fused_qkv_a_proj = DeepSeekV2FusedQkvAProjLinear(
+            config.hidden_size,
+            [self.q_lora_rank, self.kv_lora_rank + self.qk_rope_head_dim],
+            quant_config=quant_config,
+            prefix=f"{prefix}.self_attn.fused_qkv_a_proj",
+        )
+
+        # MLA Attention
+        self.attn = DeepseekV32MLAAttention(
+            vllm_config=vllm_config,
+            config=config,
+            hidden_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            qk_nope_head_dim=config.qk_nope_head_dim,
+            qk_rope_head_dim=self.qk_rope_head_dim,
+            v_head_dim=config.v_head_dim,
+            q_lora_rank=self.q_lora_rank,
+            kv_lora_rank=self.kv_lora_rank,
+            max_position_embeddings=getattr(config, "max_position_embeddings", 8192),
+            cache_config=cache_config,
+            quant_config=quant_config,
+            topk_indices_buffer=topk_indices_buffer,
+            prefix=f"{prefix}.self_attn",
+        )
+
+        # MoE or Dense MLP
+        moe_layer_freq = getattr(config, "moe_layer_freq", 1)
+        self.is_moe = (
+            config.n_routed_experts is not None
+            and layer_idx >= config.first_k_dense_replace
+            and layer_idx % moe_layer_freq == 0
+        )
+        self.routed_scaling_factor = getattr(config, "routed_scaling_factor", 1.0)
+
+        from vllm.model_executor.models.deepseek_v2 import (
+            DeepseekV2MLP,
+            DeepseekV2MoE,
+        )
+
+        if self.is_moe:
+            self.mlp = DeepseekV2MoE(
+                config=config,
+                parallel_config=parallel_config,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+        else:
+            self.mlp = DeepseekV2MLP(
+                hidden_size=config.hidden_size,
+                intermediate_size=config.intermediate_size,
+                hidden_act=config.hidden_act,
+                quant_config=quant_config,
+                prefix=f"{prefix}.mlp",
+            )
+
+    def forward(
+        self,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        residual: torch.Tensor | None,
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        if residual is None:
+            residual = hidden_states
+            hidden_states = self.input_layernorm(hidden_states)
+        else:
+            hidden_states, residual = self.input_layernorm(hidden_states, residual)
+
+        # Step 1. hidden_states -> q_c, kv_c, k_pe
+        #                       -> index_k, index_weights
+        out = self.self_attn.fused_qkv_a_proj(hidden_states)
+        if isinstance(out, tuple):
+            out = out[0]
+
+        q_c, kv_c, k_pe = out.split(
+            [self.q_lora_rank, self.kv_lora_rank, self.qk_rope_head_dim],
+            dim=-1,
+        )
+        index_k, index_weights = torch.mm(
+            hidden_states, self._fused_indexer_weights.T
+        ).split(
+            self._indexer_weights_split_sizes,
+            dim=-1,
+        )
+
+        # Steps 2-6. Combined: fused norm/rope + Q projections + sparse MLA.
+        mla = self.attn.mla_attn
+        output_shape = (hidden_states.shape[0], mla.num_heads * mla.v_head_dim)
+        output_dtype = mla.W_UV.dtype
+        attn_out = torch.empty(
+            output_shape,
+            dtype=output_dtype,
+            device=hidden_states.device,
+        )
+        attn_out = torch.ops.vllm.monolithic_attn(
+            positions,
+            q_c,
+            kv_c,
+            k_pe,
+            index_k,
+            index_weights,
+            attn_out,
+            self.layer_name,
+        )
+
+        hidden_states, _ = self.attn.o_proj(attn_out)
+        hidden_states, residual = self.post_attention_layernorm(hidden_states, residual)
+        hidden_states = self.mlp(hidden_states)
+        return hidden_states, residual
+
+    def fuse_indexer_weights(self) -> None:
+        """Fuse Step 1 and Step 3 BF16 linears used by the inlined path.
+
+        Call after model weights are loaded.
+        """
+        attn = self.attn
+        wk = attn.indexer_wk.weight.data  # [128, 7168]
+        wp = attn.indexer_weights_proj.weight.data  # [64, 7168]
+        if wk.dtype != wp.dtype:
+            raise ValueError(
+                "Cannot fuse indexer weights: expected matching dtypes for "
+                "indexer_wk and indexer_weights_proj."
+            )
+        self._fused_indexer_weights = nn.Parameter(
+            torch.cat([wk, wp], dim=0),  # [192, 7168]
+            requires_grad=False,
+        )
+        self._indexer_weights_split_sizes = [wk.shape[0], wp.shape[0]]
+
+        wq_b = attn.indexer_wq_b.weight.data
+        q_b = attn.q_b_proj.weight.data
+        if wq_b.dtype != q_b.dtype:
+            raise ValueError(
+                "Cannot fuse Step 3 weights: expected matching dtypes for "
+                "indexer_wq_b and q_b_proj."
+            )
+        self._fused_step3_q_w = nn.Parameter(
+            torch.cat([wq_b, q_b], dim=0),
+            requires_grad=False,
+        )
+        self._q_split_sizes = [wq_b.shape[0], q_b.shape[0]]
+
+
+class DeepseekV32MLAAttention(nn.Module):
+    """
+    MLA attention for DeepSeek V3.2 targeting SM100.
+    MLA forward fully inlined. MLAAttention kept only for KV cache
+    registration and backend/impl initialization.
+    """
+
+    def __init__(
+        self,
+        vllm_config: VllmConfig,
+        config,
+        hidden_size: int,
+        num_heads: int,
+        qk_nope_head_dim: int,
+        qk_rope_head_dim: int,
+        v_head_dim: int,
+        q_lora_rank: int,
+        kv_lora_rank: int,
+        max_position_embeddings: int,
+        cache_config: CacheConfig,
+        quant_config: QuantizationConfig | None,
+        topk_indices_buffer: torch.Tensor,
+        prefix: str = "",
+    ) -> None:
+        super().__init__()
+        self.hidden_size = hidden_size
+        self.qk_nope_head_dim = qk_nope_head_dim
+        self.qk_rope_head_dim = qk_rope_head_dim
+        self.qk_head_dim = qk_nope_head_dim + qk_rope_head_dim
+        self.v_head_dim = v_head_dim
+        self.q_lora_rank = q_lora_rank
+        self.kv_lora_rank = kv_lora_rank
+        self.num_heads = num_heads
+        self.num_local_heads = num_heads // get_tensor_model_parallel_world_size()
+        self.scaling = self.qk_head_dim**-0.5
+        self.rms_norm_eps = config.rms_norm_eps
+
+        # Q path
+        self.q_a_layernorm_weight = nn.Parameter(
+            torch.ones(q_lora_rank, dtype=torch.get_default_dtype())
+        )
+        self.q_b_proj = ColumnParallelLinear(
+            q_lora_rank,
+            num_heads * self.qk_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.q_b_proj",
+        )
+
+        # KV path
+        self.kv_a_layernorm_weight = nn.Parameter(
+            torch.ones(kv_lora_rank, dtype=torch.get_default_dtype())
+        )
+        self.kv_b_proj = ColumnParallelLinear(
+            kv_lora_rank,
+            num_heads * (qk_nope_head_dim + v_head_dim),
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.kv_b_proj",
+        )
+
+        # Output projection (TP sync point)
+        self.o_proj = RowParallelLinear(
+            num_heads * v_head_dim,
+            hidden_size,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.o_proj",
+        )
+
+        # RoPE
+        if config.rope_parameters["rope_type"] != "default":
+            config.rope_parameters["rope_type"] = (
+                "deepseek_yarn"
+                if config.rope_parameters.get("apply_yarn_scaling", True)
+                else "deepseek_llama_scaling"
+            )
+        self.rotary_emb = get_rope(
+            qk_rope_head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=False,
+        )
+        if config.rope_parameters["rope_type"] == "deepseek_yarn":
+            mscale_all_dim = config.rope_parameters.get("mscale_all_dim", False)
+            scaling_factor = config.rope_parameters["factor"]
+            mscale = yarn_get_mscale(scaling_factor, float(mscale_all_dim))
+            self.scaling = self.scaling * mscale * mscale
+
+        # V3.2 Sparse Indexer (inlined)
+        self.indexer_rope_emb = get_rope(
+            qk_rope_head_dim,
+            max_position=max_position_embeddings,
+            rope_parameters=config.rope_parameters,
+            is_neox_style=not getattr(config, "indexer_rope_interleave", False),
+        )
+        self.topk_tokens = config.index_topk
+        self.index_n_heads = config.index_n_heads
+        self.index_head_dim = config.index_head_dim
+        self.indexer_softmax_scale = config.index_head_dim**-0.5
+        self.indexer_quant_block_size = 128
+        self.topk_indices_buffer = topk_indices_buffer
+
+        self.indexer_wq_b = ReplicatedLinear(
+            q_lora_rank,
+            config.index_head_dim * config.index_n_heads,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.indexer.wq_b",
+        )
+        self.indexer_wk = ReplicatedLinear(
+            hidden_size,
+            config.index_head_dim,
+            bias=False,
+            quant_config=quant_config,
+            prefix=f"{prefix}.indexer.wk",
+        )
+        self.indexer_k_norm = LayerNorm(config.index_head_dim, eps=1e-6)
+        self.indexer_weights_proj = ReplicatedLinear(
+            hidden_size,
+            config.index_n_heads,
+            bias=False,
+            quant_config=None,
+            prefix=f"{prefix}.indexer.weights_proj",
+        )
+
+        idx_dim = config.index_head_dim
+        indexer_cache_head_dim = idx_dim + idx_dim // 128 * 4
+        self.indexer_k_cache = DeepseekV32IndexerCache(
+            head_dim=indexer_cache_head_dim,
+            dtype=torch.uint8,
+            prefix=f"{prefix}.indexer.k_cache",
+            cache_config=cache_config,
+        )
+        self.indexer_op = SparseAttnIndexer(
+            self.indexer_k_cache,
+            self.indexer_quant_block_size,
+            "ue8m0",
+            self.topk_tokens,
+            config.index_head_dim,
+            vllm_config.model_config.max_model_len,
+            get_max_prefill_buffer_size(vllm_config),
+            self.topk_indices_buffer,
+        )
+
+        # MLAAttention stub: only for KV cache registration + backend init.
+        # We never call its forward(); we inline everything below.
+        class _IndexerProxy:
+            def __init__(proxy_self):
+                proxy_self.topk_indices_buffer = topk_indices_buffer
+                proxy_self.indexer_op = self.indexer_op
+
+        self._indexer_proxy = _IndexerProxy()
+        self.mla_attn = MLAAttention(
+            num_heads=self.num_local_heads,
+            scale=self.scaling,
+            qk_nope_head_dim=qk_nope_head_dim,
+            qk_rope_head_dim=qk_rope_head_dim,
+            v_head_dim=v_head_dim,
+            q_lora_rank=q_lora_rank,
+            kv_lora_rank=kv_lora_rank,
+            kv_b_proj=self.kv_b_proj,
+            cache_config=cache_config,
+            quant_config=quant_config,
+            prefix=f"{prefix}.mla_attn",
+            use_sparse=True,
+            indexer=self._indexer_proxy,
+        )
+
+
+def remap_weight_name(name: str) -> str:
+    """Remap checkpoint names that differ from the module layout."""
+    replacements = [
+        (
+            "self_attn.q_a_layernorm.weight",
+            "attn.q_a_layernorm_weight",
+        ),
+        (
+            "self_attn.kv_a_layernorm.weight",
+            "attn.kv_a_layernorm_weight",
+        ),
+        ("self_attn.q_b_proj", "attn.q_b_proj"),
+        ("self_attn.kv_b_proj", "attn.kv_b_proj"),
+        ("self_attn.o_proj", "attn.o_proj"),
+        ("self_attn.indexer.", "attn.indexer_"),
+    ]
+    for old, new in replacements:
+        if old in name:
+            return name.replace(old, new)
+    return name

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/model.py

@@ -0,0 +1,151 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V3.2 NVFP4 model for SM100 (Blackwell)."""
+
+from collections.abc import Iterable
+
+import torch
+from torch import nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.distributed import get_tensor_model_parallel_world_size
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    ParallelLMHead,
+    VocabParallelEmbedding,
+)
+from vllm.platforms import current_platform
+
+from .layer import DeepseekV32DecoderLayer, remap_weight_name
+
+logger = init_logger(__name__)
+
+
+@support_torch_compile
+class DeepseekV32Model(nn.Module):
+    fall_back_to_pt_during_load = False
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.device = current_platform.device_type
+
+        topk_tokens = config.index_topk
+        self.topk_indices_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            topk_tokens,
+            dtype=torch.int32,
+            device=self.device,
+        )
+
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.embed_tokens",
+        )
+
+        self.layers = nn.ModuleList(
+            [
+                DeepseekV32DecoderLayer(
+                    vllm_config=vllm_config,
+                    config=config,
+                    layer_idx=i,
+                    topk_indices_buffer=self.topk_indices_buffer,
+                    prefix=f"{prefix}.layers.{i}",
+                )
+                for i in range(config.num_hidden_layers)
+            ]
+        )
+
+        self.norm = RMSNorm(
+            hidden_size=config.hidden_size,
+            eps=config.rms_norm_eps,
+            dtype=torch.get_default_dtype(),
+        )
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+    ) -> torch.Tensor:
+        hidden_states = self.embed_tokens(input_ids)
+        residual = None
+        for layer in self.layers:
+            hidden_states, residual = layer(positions, hidden_states, residual)
+        hidden_states, _ = self.norm(hidden_states, residual)
+        return hidden_states
+
+
+class DeepseekV32ForCausalLM(nn.Module):
+    packed_modules_mapping = {
+        "gate_up_proj": ["gate_proj", "up_proj"],
+        "fused_qkv_a_proj": ["q_a_proj", "kv_a_proj_with_mqa"],
+    }
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__()
+        config = vllm_config.model_config.hf_config
+        quant_config = vllm_config.quant_config
+        self.config = config
+        self.quant_config = quant_config
+        self.tp_size = get_tensor_model_parallel_world_size()
+
+        self.model = DeepseekV32Model(
+            vllm_config=vllm_config,
+            prefix=f"{prefix}.model" if prefix else "model",
+        )
+        self.lm_head = ParallelLMHead(
+            config.vocab_size,
+            config.hidden_size,
+            quant_config=quant_config,
+            prefix=f"{prefix}.lm_head" if prefix else "lm_head",
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+        self.num_redundant_experts = 0
+
+    def embed_input_ids(self, input_ids: torch.Tensor) -> torch.Tensor:
+        return self.model.embed_tokens(input_ids)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor,
+        positions: torch.Tensor,
+        intermediate_tensors=None,
+        inputs_embeds=None,
+    ) -> torch.Tensor:
+        return self.model(input_ids, positions)
+
+    def compute_logits(self, hidden_states: torch.Tensor) -> torch.Tensor | None:
+        return self.logits_processor(self.lm_head, hidden_states)
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        """Delegate to the original DeepSeek V2 weight loader.
+
+        Our module structure matches the original for all weights that
+        need special loading (fused_qkv_a_proj, experts, gate_up_proj).
+        Only layernorm weights and indexer paths differ.
+        """
+        from vllm.model_executor.models.deepseek_v2 import (
+            DeepseekV2ForCausalLM,
+        )
+
+        def _remap_weights():
+            for name, w in weights:
+                yield remap_weight_name(name), w
+
+        self.use_mha = False
+        self.fuse_qkv_a_proj = True
+        self.is_fp4_ckpt = False
+        loaded = DeepseekV2ForCausalLM.load_weights(self, _remap_weights())
+
+        # Fuse indexer linear weights after loading.
+        for layer in self.model.layers:
+            layer.fuse_indexer_weights()
+
+        return loaded

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/mtp.py

@@ -0,0 +1,168 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""DeepSeek V3.2 MTP model for SM100 (Blackwell)."""
+
+from collections.abc import Iterable
+
+import torch
+import torch.nn as nn
+
+from vllm.compilation.decorators import support_torch_compile
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.layers.layernorm import RMSNorm
+from vllm.model_executor.layers.logits_processor import LogitsProcessor
+from vllm.model_executor.layers.vocab_parallel_embedding import (
+    VocabParallelEmbedding,
+)
+from vllm.model_executor.models.deepseek_mtp import DeepSeekMTP as DeepSeekMTPBase
+from vllm.model_executor.models.deepseek_mtp import (
+    DeepSeekMultiTokenPredictor as DeepSeekMultiTokenPredictorBase,
+)
+from vllm.model_executor.models.deepseek_mtp import (
+    DeepSeekMultiTokenPredictorLayer as DeepSeekMultiTokenPredictorLayerBase,
+)
+from vllm.model_executor.models.deepseek_mtp import SharedHead as SharedHeadBase
+from vllm.model_executor.models.deepseek_v2 import DeepseekV2MoE
+from vllm.model_executor.models.utils import maybe_prefix
+from vllm.platforms import current_platform
+from vllm.sequence import IntermediateTensors
+
+from .layer import DeepseekV32DecoderLayer
+from .model import remap_weight_name
+
+logger = init_logger(__name__)
+
+
+class SharedHead(SharedHeadBase):
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return rms_norm(hidden_states, self.norm.weight, self.norm.variance_epsilon)
+
+
+class DeepSeekMultiTokenPredictorLayer(DeepSeekMultiTokenPredictorLayerBase):
+    def __init__(self, vllm_config: VllmConfig, prefix: str) -> None:
+        nn.Module.__init__(self)
+
+        assert vllm_config.speculative_config is not None
+        config = vllm_config.speculative_config.draft_model_config.hf_config
+        quant_config = vllm_config.quant_config
+
+        self.config = config
+        self.enorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.hnorm = RMSNorm(config.hidden_size, eps=config.rms_norm_eps)
+        self.eh_proj = nn.Linear(config.hidden_size * 2, config.hidden_size, bias=False)
+
+        topk_indices_buffer = torch.empty(
+            vllm_config.scheduler_config.max_num_batched_tokens,
+            config.index_topk,
+            dtype=torch.int32,
+            device=current_platform.device_type,
+        )
+
+        self.shared_head = SharedHead(
+            config=config, prefix=prefix, quant_config=quant_config
+        )
+        self.mtp_block = DeepseekV32DecoderLayer(
+            vllm_config=vllm_config,
+            config=config,
+            layer_idx=int(prefix.rsplit(".", 1)[-1]),
+            topk_indices_buffer=topk_indices_buffer,
+            prefix=prefix,
+        )
+
+
+class DeepSeekMultiTokenPredictor(DeepSeekMultiTokenPredictorBase):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        config = vllm_config.model_config.hf_config
+        self.mtp_start_layer_idx = config.num_hidden_layers
+        self.num_mtp_layers = config.num_nextn_predict_layers
+
+        self.layers = torch.nn.ModuleDict(
+            {
+                str(idx): DeepSeekMultiTokenPredictorLayer(
+                    vllm_config, f"{prefix}.layers.{idx}"
+                )
+                for idx in range(
+                    self.mtp_start_layer_idx,
+                    self.mtp_start_layer_idx + self.num_mtp_layers,
+                )
+            }
+        )
+        self.embed_tokens = VocabParallelEmbedding(
+            config.vocab_size,
+            config.hidden_size,
+            prefix=maybe_prefix(prefix, "embed_tokens"),
+        )
+        self.logits_processor = LogitsProcessor(config.vocab_size)
+
+
+@support_torch_compile
+class DeepSeekMTP(DeepSeekMTPBase):
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
+        nn.Module.__init__(self)
+        self.config = vllm_config.model_config.hf_config
+        self.quant_config = vllm_config.quant_config
+        assert hasattr(self.config, "index_topk")
+        cache_config = vllm_config.cache_config
+        if cache_config.cache_dtype == "bfloat16":
+            cache_config.cache_dtype = "auto"
+            logger.info("Using bfloat16 kv-cache for DeepSeekV3.2")
+        self.model = DeepSeekMultiTokenPredictor(
+            vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")
+        )
+        self.set_moe_parameters()
+        # Keep the original loader from applying the fused FP4 indexer remap.
+        self.is_fp4_ckpt = False
+
+    def set_moe_parameters(self):
+        self.expert_weights = []
+        self.num_moe_layers = self.config.num_nextn_predict_layers
+        self.num_expert_groups = self.config.n_group
+
+        self.moe_layers = []
+        self.moe_mlp_layers = []
+        example_moe = None
+        for layer in self.model.layers.values():
+            layer = layer.mtp_block
+            assert isinstance(layer, DeepseekV32DecoderLayer)
+            if isinstance(layer.mlp, DeepseekV2MoE):
+                example_moe = layer.mlp
+                self.moe_mlp_layers.append(layer.mlp)
+                self.moe_layers.append(layer.mlp.experts)
+        self.extract_moe_parameters(example_moe)
+
+    def forward(
+        self,
+        input_ids: torch.Tensor | None,
+        positions: torch.Tensor,
+        hidden_states: torch.Tensor,
+        intermediate_tensors: IntermediateTensors | None = None,
+        inputs_embeds: torch.Tensor | None = None,
+        spec_step_idx: int = 0,
+    ) -> torch.Tensor:
+        del intermediate_tensors
+        return self.model(
+            input_ids, positions, hidden_states, inputs_embeds, spec_step_idx
+        )
+
+    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
+        loaded_params = super().load_weights(weights)
+        for layer in self.model.layers.values():
+            layer.mtp_block.fuse_indexer_weights()
+        return loaded_params
+
+    def _rewrite_spec_layer_name(self, spec_layer: int, name: str) -> str:
+        name = super()._rewrite_spec_layer_name(spec_layer, name)
+        return remap_weight_name(name)
+
+
+@torch.compile
+def rms_norm(x: torch.Tensor, w: torch.Tensor, eps: float) -> torch.Tensor:
+    orig_dtype = x.dtype
+    x = x.to(torch.float32)
+    mean_sq = (x * x).mean(dim=-1, keepdim=True)
+    rrms = torch.rsqrt(mean_sq + eps)
+    x = x * rrms
+    x = x * w.to(torch.float32)
+    return x.to(orig_dtype)

vllm/model_executor/specialized_models/deepseek_v3_2_nvfp4/sparse_indexer.py

@@ -0,0 +1,175 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Custom Sparse Attention Indexer layers."""
+
+import torch
+
+import vllm.envs as envs
+from vllm import _custom_ops as ops
+from vllm.forward_context import get_forward_context
+from vllm.logger import init_logger
+from vllm.platforms import current_platform
+from vllm.utils.deep_gemm import fp8_mqa_logits, fp8_paged_mqa_logits
+from vllm.utils.torch_utils import (
+    LayerNameType,
+    _resolve_layer_name,
+)
+from vllm.v1.attention.backends.mla.indexer import (
+    DeepseekV32IndexerMetadata,
+)
+from vllm.v1.attention.ops.common import pack_seq_triton, unpack_seq_triton
+from vllm.v1.worker.workspace import current_workspace_manager
+
+logger = init_logger(__name__)
+
+RADIX_TOPK_WORKSPACE_SIZE = 1024 * 1024
+
+
+def sparse_attn_indexer(
+    k_cache_prefix: LayerNameType,
+    kv_cache: torch.Tensor,
+    q_fp8: torch.Tensor,
+    weights: torch.Tensor,
+    topk_tokens: int,
+    head_dim: int,
+    max_model_len: int,
+    total_seq_lens: int,
+    topk_indices_buffer: torch.Tensor,
+) -> torch.Tensor:
+    # careful! this will be None in dummy run
+    attn_metadata = get_forward_context().attn_metadata
+    fp8_dtype = current_platform.fp8_dtype()
+    k_cache_prefix = _resolve_layer_name(k_cache_prefix)
+
+    # assert isinstance(attn_metadata, dict)
+    if not isinstance(attn_metadata, dict):
+        # Reserve workspace for indexer during profiling run
+        current_workspace_manager().get_simultaneous(
+            ((total_seq_lens, head_dim), torch.float8_e4m3fn),
+            ((total_seq_lens, 4), torch.uint8),
+            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+        )
+
+        # Dummy allocation to simulate for peak logits tensor memory during inference.
+        # FP8 elements so elements == bytes
+        max_logits_elems = envs.VLLM_SPARSE_INDEXER_MAX_LOGITS_MB * 1024 * 1024
+        _ = torch.empty(max_logits_elems, dtype=torch.uint8, device=q_fp8.device)
+        return None
+
+    attn_metadata = attn_metadata[k_cache_prefix]  # type: ignore[assignment]
+    assert isinstance(attn_metadata, DeepseekV32IndexerMetadata)
+    has_decode = attn_metadata.num_decodes > 0
+    has_prefill = attn_metadata.num_prefills > 0
+    num_decode_tokens = attn_metadata.num_decode_tokens
+
+    if has_prefill:
+        prefill_metadata = attn_metadata.prefill
+        assert prefill_metadata is not None
+
+        # Get the full shared workspace buffers once (will allocate on first use)
+        workspace_manager = current_workspace_manager()
+        k_fp8_full, k_scale_full = workspace_manager.get_simultaneous(
+            ((total_seq_lens, head_dim), fp8_dtype),
+            ((total_seq_lens, 4), torch.uint8),
+        )
+        for chunk in prefill_metadata.chunks:
+            k_fp8 = k_fp8_full[: chunk.total_seq_lens]
+            k_scale = k_scale_full[: chunk.total_seq_lens]
+
+            if not chunk.skip_kv_gather:
+                ops.cp_gather_indexer_k_quant_cache(
+                    kv_cache,
+                    k_fp8,
+                    k_scale,
+                    chunk.block_table,
+                    chunk.cu_seq_lens,
+                )
+
+            logits = fp8_mqa_logits(
+                q_fp8[chunk.token_start : chunk.token_end],
+                (k_fp8, k_scale.view(torch.float32).flatten()),
+                weights[chunk.token_start : chunk.token_end],
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+                clean_logits=False,
+            )
+            num_rows = logits.shape[0]
+
+            topk_indices = topk_indices_buffer[
+                chunk.token_start : chunk.token_end, :topk_tokens
+            ]
+
+            torch.ops._C.top_k_per_row_prefill(
+                logits,
+                chunk.cu_seqlen_ks,
+                chunk.cu_seqlen_ke,
+                topk_indices,
+                num_rows,
+                logits.stride(0),
+                logits.stride(1),
+                topk_tokens,
+            )
+
+    if has_decode:
+        decode_metadata = attn_metadata.decode
+        assert decode_metadata is not None
+        # kv_cache shape [
+        # kv_cache size requirement [num_block, block_size, n_head, head_dim],
+        # we only have [num_block, block_size, head_dim],
+        kv_cache = kv_cache.unsqueeze(-2)
+        decode_lens = decode_metadata.decode_lens
+        if decode_metadata.requires_padding:
+            # pad in edge case where we have short chunked prefill length <
+            # decode_threshold since we unstrictly split
+            # prefill and decode by decode_threshold
+            # (currently set to 1 + speculative tokens)
+            padded_q_fp8_decode_tokens = pack_seq_triton(
+                q_fp8[:num_decode_tokens], decode_lens
+            )
+        else:
+            padded_q_fp8_decode_tokens = q_fp8[:num_decode_tokens].reshape(
+                decode_lens.shape[0], -1, *q_fp8.shape[1:]
+            )
+        # TODO: move and optimize below logic with triton kernels
+        batch_size = padded_q_fp8_decode_tokens.shape[0]
+        next_n = padded_q_fp8_decode_tokens.shape[1]
+        num_padded_tokens = batch_size * next_n
+        seq_lens = decode_metadata.seq_lens[:batch_size]
+        # seq_lens is (B, next_n) for native spec decode, (B,) otherwise.
+        # fp8_paged_mqa_logits and all topk kernels accept both shapes.
+        logits = fp8_paged_mqa_logits(
+            padded_q_fp8_decode_tokens,
+            kv_cache,
+            weights[:num_padded_tokens],
+            seq_lens,
+            decode_metadata.block_table,
+            decode_metadata.schedule_metadata,
+            max_model_len=max_model_len,
+            clean_logits=False,
+        )
+        num_rows = logits.shape[0]
+        topk_indices = topk_indices_buffer[:num_padded_tokens, :topk_tokens]
+
+        workspace_manager = current_workspace_manager()
+        (topk_workspace,) = workspace_manager.get_simultaneous(
+            ((RADIX_TOPK_WORKSPACE_SIZE,), torch.uint8),
+        )
+        torch.ops._C.persistent_topk(
+            logits,
+            seq_lens,
+            topk_indices,
+            topk_workspace,
+            topk_tokens,
+            attn_metadata.max_seq_len,
+        )
+
+        if decode_metadata.requires_padding:
+            # if padded, we need to unpack
+            # the topk indices removing padded tokens
+            topk_indices = unpack_seq_triton(
+                topk_indices.reshape(batch_size, -1, topk_indices.shape[-1]),
+                decode_lens,
+            )
+            topk_indices_buffer[: topk_indices.shape[0], : topk_indices.shape[-1]] = (
+                topk_indices
+            )

vllm/v1/worker/gpu/cudagraph_utils.py

@@ -335,6 +335,7 @@ class ModelCudaGraphManager(CudaGraphManager):
                 with set_forward_context(
                     attn_metadata if cg_mode != CUDAGraphMode.PIECEWISE else None,
                     self.vllm_config,
+                    num_actual_tokens=input_buffers.num_actual_tokens,
                     num_tokens=num_tokens,
                     cudagraph_runtime_mode=cg_mode,
                     num_tokens_across_dp=num_tokens_across_dp,

vllm/v1/worker/gpu/input_batch.py

@@ -30,6 +30,7 @@ class InputBuffers:
         self.dcp_local_seq_lens = torch.zeros(
             max_num_reqs, dtype=torch.int32, device=device
         )
+        self.num_actual_tokens = torch.zeros(1, dtype=torch.int32, device=device)
 
 
 @dataclass

vllm/v1/worker/gpu/model_runner.py

@@ -971,6 +971,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             empty_output = self.kv_connector.no_forward(scheduler_output)
             return empty_output
 
+        self.input_buffers.num_actual_tokens[:1] = num_toks if not dummy_run else 0
         if not dummy_run:
             # Common case.
             # Prepare all the inputs and copy to the input buffers.
@@ -1075,6 +1076,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
             with set_forward_context(
                 attn_metadata,
                 self.vllm_config,
+                num_actual_tokens=self.input_buffers.num_actual_tokens,
                 num_tokens=input_batch.num_tokens_after_padding,
                 cudagraph_runtime_mode=batch_desc.cg_mode,
                 num_tokens_across_dp=num_tokens_across_dp,