@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 9163190dda

[ROCm][Bugfix][Perf] Tune multi-stream shared experts use; wvSplitKrc fixes (#56098)

Signed-off-by: Matthew Wong <Matthew.Wong2@amd.com> Co-authored-by: Claude Opus 5 (1M context) <noreply@anthropic.com>

Matt committed Sep 10, 2026 at 15:14 UTC 9163190dda009a310d8c175d63860ac7c671ca90
6 files changed +369 -77
csrc/rocm/skinny_gemms.cu
+216 -32
@@ -8,12 +8,20 @@
8
9 #include <stdexcept>
10 #include <algorithm>
11 +#include <map>
12 +#include <mutex>
13 +#include <utility>
14
15 #include "../cuda_compat.h"
16 #include "dispatch_utils.h"
17 #include "quantization/w8a8/fp8/common.cuh"
18 #include "core/batch_invariant.hpp"
19
20 +// Number of streams the pre-allocated split-K pool covers. ~7.5 MiB per slot.
21 +static constexpr int64_t kWvSlots = 8;
22 +// Slots held back for capturing streams, which cannot allocate one themselves.
23 +static constexpr int64_t kWvCaptureSlots = 8;
24 +
25 // TODO(rasmith): The kernels in this file are susceptible to integer overflow
26 // issues, do not take strides, and are unable to handle PyTorch tensors that
27 // return is_contiguous() as False (the tensors may actually be contiguous
@@ -1364,9 +1372,18 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1372 };
1373 using big4 = __attribute__((__vector_size__(4 * sizeof(bigType)))) __bf16;
1374
1367 - __shared__ scalar_t stg[WvPrGrp * WVLDS / GrpsShrB];
1375 + // The main-loop staging areas are separated from the split-K readback by a
1376 + // __syncthreads(), so the readback may reuse all of LDS.
1377 + __shared__ union {
1378 + struct {
1379 + scalar_t stg[WvPrGrp * WVLDS / GrpsShrB];
1380 + scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB];
1381 + } mainloop;
1382 + scalar_t all[max_lds_len];
1383 + } lds;
1384 + auto& stg = lds.mainloop.stg;
1385 + auto& s = lds.mainloop.s;
1386 unsigned int* myStg = (unsigned int*)(&stg[WVLDS * (threadIdx.y / GrpsShrB)]);
1369 - __shared__ scalar_t s[max_lds_len - WvPrGrp * WVLDS / GrpsShrB];
1387
1388 #ifndef WVSPLITKRC_1KPASS
1389 constexpr int TUC_ = (THRDS * UNRL * A_CHUNK);
@@ -1682,33 +1699,69 @@ __global__ void __launch_bounds__(WvPrGrp* THRDS)
1699 if (my_cntr + 1 == k_rnd) {
1700 cntr[adr_] = 0; // clear for next round
1701 if constexpr (DTRMNSTC) {
1702 + constexpr uint32_t NT = N / NTILE / GrpsShrB;
1703 + // The readback stages k-shard partials through LDS, which bounds how
1704 + // many shards may be in flight at once.
1705 + constexpr uint32_t LDS_F4 =
1706 + max_lds_len * sizeof(scalar_t) / sizeof(float4);
1707 + constexpr uint32_t KSBMAX = LDS_F4 / (THRDS * 4 * NT);
1708 + static_assert(KSBMAX >= 1, "LDS cannot stage even one k-shard");
1709 + const uint32_t KSB = (k_rnd < KSBMAX) ? k_rnd : KSBMAX;
1710 +
1711 + // First batch is peeled; it is the only batch for most shapes.
1712 #pragma unroll
1686 - for (int ks = 0; ks < k_rnd; ks++) {
1687 - for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1713 + for (uint32_t kb = 0; kb < KSB; kb++) {
1714 + for (uint32_t nt = 0; nt < NT; nt++) {
1715 int g_nindx =
1716 (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1717 int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1718 __builtin_amdgcn_global_load_lds(
1692 - (float4*)(&glbl[g_adr + M * N * ks]),
1693 - &(((float4*)s)[(threadIdx.y * THRDS) + ks * THRDS * 4 +
1694 - nt * THRDS * 4 * k_rnd]),
1719 + (float4*)(&glbl[g_adr + M * N * kb]),
1720 + &(((float4*)lds.all)[(threadIdx.y * THRDS) + kb * THRDS * 4 +
1721 + nt * THRDS * 4 * KSB]),
1722 16, 0, 0);
1723 }
1724 }
1698 - if (BIAS)
1699 - for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1725 + if (BIAS) // overlaps the DMAs above
1726 + for (uint32_t nt = 0; nt < NT; nt++) {
1727 for (uint32_t j = 0; j < 4; j++) {
1728 int nindx = (j + (threadIdx.x / 16) * 4) + nt * NTILE +
1729 (N / GrpsShrB) * (threadIdx.y % GrpsShrB);
1730 biases[nt][j] = BIAS[(mindx % Bx) + (nindx % By) * Bx];
1731 }
1732 }
1706 - asm volatile("s_waitcnt 0");
1707 - for (int ks = 0; ks < k_rnd; ks++) {
1708 - for (uint32_t nt = 0; nt < N / NTILE / GrpsShrB; nt++) {
1709 - float4 eval = ((float4*)s)[(threadIdx.x + threadIdx.y * THRDS) +
1710 - ks * THRDS * 4 + nt * THRDS * 4 * k_rnd];
1711 - vals[nt].f4 += eval;
1733 + asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
1734 + for (uint32_t kb = 0; kb < KSB; kb++) {
1735 + for (uint32_t nt = 0; nt < NT; nt++) {
1736 + vals[nt].f4 +=
1737 + ((float4*)lds.all)[(threadIdx.x + threadIdx.y * THRDS) +
1738 + kb * THRDS * 4 + nt * THRDS * 4 * KSB];
1739 + }
1740 + }
1741 +
1742 + for (uint32_t ks0 = KSB; ks0 < k_rnd; ks0 += KSB) {
1743 + const uint32_t ksn = (k_rnd - ks0 < KSB) ? (k_rnd - ks0) : KSB;
1744 + // prior reads must retire before the staging is overwritten
1745 + asm volatile("s_waitcnt lgkmcnt(0)" ::: "memory");
1746 + for (uint32_t kb = 0; kb < ksn; kb++) {
1747 + for (uint32_t nt = 0; nt < NT; nt++) {
1748 + int g_nindx =
1749 + (nt * NTILE + (N / GrpsShrB) * (threadIdx.y % GrpsShrB)) / 4;
1750 + int g_adr = g_mindx * 4 + 0 + M * g_nindx * 4;
1751 + __builtin_amdgcn_global_load_lds(
1752 + (float4*)(&glbl[g_adr + M * N * (ks0 + kb)]),
1753 + &(((float4*)lds.all)[(threadIdx.y * THRDS) + kb * THRDS * 4 +
1754 + nt * THRDS * 4 * KSB]),
1755 + 16, 0, 0);
1756 + }
1757 + }
1758 + asm volatile("s_waitcnt vmcnt(0)" ::: "memory");
1759 + for (uint32_t kb = 0; kb < ksn; kb++) {
1760 + for (uint32_t nt = 0; nt < NT; nt++) {
1761 + vals[nt].f4 +=
1762 + ((float4*)lds.all)[(threadIdx.x + threadIdx.y * THRDS) +
1763 + kb * THRDS * 4 + nt * THRDS * 4 * KSB];
1764 + }
1765 }
1766 }
1767 } else {
@@ -1771,7 +1824,7 @@ __global__ void wvSplitKrc_(const int actlN, const int K, const int Kap,
1824 torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1825 const std::optional<at::Tensor>& in_bias,
1826 const int64_t CuCount) {
1774 - int _DTRMNSTC = 1; // vllm::vllm_is_batch_invariant();
1827 + static constexpr bool _DTRMNSTC = true;
1828
1829 auto M_in = in_b.size(0);
1830 auto N_in = in_a.size(0);
@@ -1819,26 +1872,157 @@ torch::Tensor wvSplitKrc(const at::Tensor& in_a, const at::Tensor& in_b,
1872
1873 if (CuNeeded > CuCount) throw std::runtime_error("Invalid wvSplitKrc size");
1874
1822 - // Can we increase SplitK by shrinking the K-shared to 256?
1823 - int chunkk = (CuNeeded * 2 <= CuCount) ? 2 : 1;
1824 -
1825 - static torch::Tensor axl_glbl =
1826 - torch::zeros(
1827 - 128 * 1024 * (_DTRMNSTC ? 12 : 1),
1828 - torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()))
1829 - .detach();
1830 - static torch::Tensor axl_cntr =
1831 - torch::zeros(
1832 - 128 * 1024 * (_DTRMNSTC ? 12 : 1) / 4,
1833 - torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
1834 - .detach();
1835 - auto glbl = axl_glbl.data_ptr<float>();
1836 - auto cntr = axl_cntr.data_ptr<int>();
1875 + // Can we increase SplitK by shrinking the K-shard to 256? That doubles both
1876 + // the CUs needed and the k-shards to read back, and the readback is serial in
1877 + // the shard count, so it pays only up to a crossover. The crossover is not a
1878 + // constant -- it falls as M grows -- so one cap cannot sit on it for every
1879 + // shape; this value minimizes total time over the admitted shapes. Retune it
1880 + // if the readback cost changes, and note the directions are not symmetric:
1881 + // capping too high costs about twice what capping too low does.
1882 + // Mirrored by rocm_unquantized_gemm_impl() in
1883 + // vllm/model_executor/layers/utils.py; both must pick the same chunkk or the
1884 + // host-side fit check bounds the wrong k_rnd.
1885 + constexpr int64_t CHUNKK2_MAX_SHARDS = 11;
1886 + const int64_t shards_chunkk2 = (K_in + 256 - 1) / 256; // 256 == 512 / 2
1887 + const int chunkk = (N_p2 != 16 && CuNeeded * 2 <= CuCount &&
1888 + shards_chunkk2 <= CHUNKK2_MAX_SHARDS)
1889 + ? 2
1890 + : 1;
1891 +
1892 + // One fp32 partial per (M, N, k-shard) must fit the split-K workspace.
1893 + const int64_t k_rnd = (K_in + 512 / chunkk - 1) / (512 / chunkk);
1894 + constexpr int64_t glbl_cap = 128 * 1024 * (_DTRMNSTC ? 12 : 1);
1895 + const int64_t glbl_needed =
1896 + _DTRMNSTC ? (int64_t)N_p2 * M_in * k_rnd : (int64_t)N_p2 * M_in;
1897 + TORCH_CHECK(glbl_needed <= glbl_cap,
1898 + "wvSplitKrc: split-K workspace too small", " (need ", glbl_needed,
1899 + " > ", glbl_cap, " floats) for N=", N_in, " K=", K_in,
1900 + " M=", M_in);
1901 +
1902 + // Split-K workspaces: one pre-partitioned pool per device.
1903 + // These must be zeroed before first use, and the kernel re-zeros them after
1904 + // each use.
1905 + // To avoid unnecessary overhead, we do not allocate them zeroed
1906 + // per-invocation. However, a simple static allocation does not work with
1907 + // multi-streams, because concurrent streams can alias into the same
1908 + // workspace. Solution: pre-allocate kWvSlots = 8 slots per device, plus
1909 + // kWvCaptureSlots reserved for capturing streams, and zero the whole pool
1910 + // once at creation. On first use, a stream takes one of the slots if any are
1911 + // still available. On subsequent uses, it re-uses the slot. If no slots are
1912 + // available, fall back to allocating a zeroed one per call, which is safe but
1913 + // adds overhead.
1914 + // A capturing stream must neither allocate nor enqueue anything -- either
1915 + // would be captured into the graph -- so it is only handed a reserved slot.
1916 + struct WvSplitKrcPool {
1917 + torch::Tensor glbl, cntr;
1918 + int64_t glbl_stride = 0, cntr_stride = 0;
1919 + int next_slot = 0;
1920 + int next_capture_slot = 0;
1921 + std::map<cudaStream_t, int> slot_of;
1922 + std::map<cudaStream_t, int> capture_slot_of;
1923 + };
1924 + // mutex not strictly necessary if caller is always single-threaded (Python)
1925 + // but this preserves Torch thread_local semantics for the C++ API
1926 + static std::mutex wv_pool_mu;
1927 + static std::map<int, WvSplitKrcPool> wv_pools;
1928 +
1929 + float* glbl;
1930 + int* cntr;
1931 + // Fallback space
1932 + torch::Tensor spill_glbl, spill_cntr;
1933 + {
1934 + std::lock_guard<std::mutex> wv_lk(wv_pool_mu);
1935 + const int wv_dev = static_cast<int>(in_a.device().index());
1936 + WvSplitKrcPool& P = wv_pools[wv_dev];
1937 + constexpr int64_t wv_gf = 128 * 1024 * (_DTRMNSTC ? 12 : 1);
1938 + constexpr int64_t wv_ci = wv_gf / 4;
1939 + constexpr int64_t wv_all_slots = kWvSlots + kWvCaptureSlots;
1940 +
1941 + // A stream whose capture was already invalidated reports an error; treat
1942 + // that as not capturing, since the capture is lost either way.
1943 + cudaStreamCaptureStatus capture_status = cudaStreamCaptureStatusNone;
1944 + unsigned long long capture_id = 0;
1945 + if (cudaStreamGetCaptureInfo(stream, &capture_status, &capture_id) !=
1946 + cudaSuccess) {
1947 + capture_status = cudaStreamCaptureStatusNone;
1948 + }
1949 + const bool wv_capturing = capture_status != cudaStreamCaptureStatusNone;
1950 +
1951 + if (!P.glbl.defined()) {
1952 + TORCH_CHECK(!wv_capturing,
1953 + "wvSplitKrc: the split-K workspace pool cannot be allocated "
1954 + "during graph capture. Call "
1955 + "warmup_rocm_skinny_gemm_workspaces() "
1956 + "(vllm/model_executor/layers/utils.py) eagerly first.");
1957 + P.glbl_stride = wv_gf;
1958 + P.cntr_stride = wv_ci;
1959 + P.glbl = torch::zeros(wv_gf * wv_all_slots, torch::TensorOptions()
1960 + .dtype(torch::kFloat32)
1961 + .device(in_a.device()))
1962 + .detach();
1963 + P.cntr =
1964 + torch::zeros(
1965 + wv_ci * wv_all_slots,
1966 + torch::TensorOptions().dtype(torch::kInt).device(in_a.device()))
1967 + .detach();
1968 + // The mutex orders host access to the map, not device execution, so it
1969 + // alone would let another stream be handed a slot the fill above is
1970 + // still clearing. Synchronizing here, inside the lock, is what makes
1971 + // every slot safe to use without any further per-slot zeroing.
1972 + const cudaError_t wv_err = cudaStreamSynchronize(stream);
1973 + TORCH_CHECK(wv_err == cudaSuccess,
1974 + "wvSplitKrc: failed to zero the split-K workspace pool: ",
1975 + cudaGetErrorString(wv_err));
1976 + }
1977 +
1978 + int wv_slot = -1;
1979 + if (wv_capturing) {
1980 + // One workspace per capturing stream, so graphs captured on the same
1981 + // stream share it -- safe because at most one of them replays at a time
1982 + // (vLLM captures one graph per padded batch size, and runs one per step).
1983 + auto cap_it = P.capture_slot_of.find(stream);
1984 + if (cap_it == P.capture_slot_of.end()) {
1985 + TORCH_CHECK(P.next_capture_slot < kWvCaptureSlots,
1986 + "wvSplitKrc: no free capture workspace slot; more than ",
1987 + kWvCaptureSlots,
1988 + " distinct streams have captured a graph containing "
1989 + "wvSplitKrc. Raise kWvCaptureSlots in "
1990 + "csrc/rocm/skinny_gemms.cu.");
1991 + cap_it = P.capture_slot_of
1992 + .emplace(stream, static_cast<int>(kWvSlots) +
1993 + P.next_capture_slot++)
1994 + .first;
1995 + }
1996 + wv_slot = cap_it->second;
1997 + } else {
1998 + auto slot_it = P.slot_of.find(stream);
1999 + if (slot_it != P.slot_of.end()) {
2000 + wv_slot = slot_it->second;
2001 + } else if (P.next_slot < kWvSlots) {
2002 + wv_slot = P.next_slot++;
2003 + P.slot_of.emplace(stream, wv_slot);
2004 + }
2005 + }
2006 +
2007 + if (wv_slot >= 0) {
2008 + glbl = P.glbl.data_ptr<float>() + wv_slot * P.glbl_stride;
2009 + cntr = P.cntr.data_ptr<int>() + wv_slot * P.cntr_stride;
2010 + } else {
2011 + spill_glbl = torch::zeros(
2012 + wv_gf,
2013 + torch::TensorOptions().dtype(torch::kFloat32).device(in_a.device()));
2014 + spill_cntr = torch::zeros(
2015 + wv_ci,
2016 + torch::TensorOptions().dtype(torch::kInt).device(in_a.device()));
2017 + glbl = spill_glbl.data_ptr<float>();
2018 + cntr = spill_cntr.data_ptr<int>();
2019 + }
2020 + }
2021
2022 #define WVSPLITKrc(_N, _GrpsShrB, _CHUNKK) \
2023 { \
2024 dim3 block(64, 4); \
1841 - if (_DTRMNSTC) \
2025 + if constexpr (_DTRMNSTC) \
2026 wvSplitKrc_<fptype, 64, 16, 4, 8, 1, _N, _GrpsShrB, _CHUNKK, 1> \
2027 <<<grid, block, 0, stream>>>(N_in, K_in, Kap_in, M_in, Bx_in, By_in, \
2028 af4, bf4, biasf4, glbl, cntr, c, \
tests/kernels/quantization/test_rocm_skinny_gemms.py
+97
@@ -53,6 +53,17 @@ OPTIONS_WVSPLITK_FP8 = [
53
54 DTYPES = [torch.bfloat16, torch.float16]
55
56 +# kWvSlots in csrc/rocm/skinny_gemms.cu; beyond it streams take overflow
57 +# workspaces.
58 +WVSPLITKRC_SLOTS = 8
59 +
60 +# The CU count is the grid size. Production passes num_compute_units(); a small
61 +# value manufactures better concurrency for these tests.
62 +WVSPLITKRC_TEST_CU = 16
63 +
64 +# Gate spin, ~125 ms on gfx950 -- must outlast the host enqueue loop it gates.
65 +_GATE_CYCLES = 300_000_000
66 +
67 # Specific (N, K, M) combinations for targeted testing
68 NKM_FACTORS_LLMM1 = [
69 # Small, medium, large cases
@@ -116,6 +127,15 @@ K_FACTORS_WVSPLITKRC = [2880, 2880 + 8, 3072, 3072 + 8]
127 # M tiles are 64 rows, +16 for a partial tile
128 M_FACTORS_WVSPLITKRC = [128, 128 + 16, 256, 256 + 16, 640, 640 + 16]
129
130 +# (N, K, M) with more K-shards than the readback can stage in one LDS pass.
131 +# Spans both CHUNKK values and both N-tile counts.
132 +NKM_FACTORS_WVSPLITKRC_LARGE_K = [
133 + (128, 6144, 128),
134 + (96, 8192, 128),
135 + (128, 12288, 128),
136 + (32, 12288, 128),
137 +]
138 +
139 NKM_FACTORS_WVSPLITK_FP8 = [
140 # FP8-specific cases with K % 16 == 0
141 (1, 16, 16),
@@ -222,6 +242,83 @@ def test_rocm_wvsplitkrc_kernel(n, k, m, dtype, padded_a, bias_mode, xnorm, seed
242 torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
243
244
245 +@pytest.mark.parametrize("n,k,m", NKM_FACTORS_WVSPLITKRC_LARGE_K)
246 +@pytest.mark.parametrize("dtype", DTYPES)
247 +@pytest.mark.parametrize("seed", SEEDS)
248 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
249 +@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
250 +def test_rocm_wvsplitkrc_large_k(n, k, m, dtype, seed):
251 + """K large enough that the split-K readback must stage LDS in batches."""
252 + torch.manual_seed(seed)
253 +
254 + xavier = math.sqrt(2 / k)
255 + A = (torch.rand(n, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
256 + B = (torch.rand(m, k, dtype=dtype, device="cuda") * 2 - 1) * xavier
257 +
258 + ref_out = torch.nn.functional.linear(A, B, None)
259 + out = ops.wvSplitKrc(A, B, num_compute_units(), None)
260 +
261 + torch.testing.assert_close(out, ref_out, atol=1e-3, rtol=1e-2)
262 +
263 +
264 +@pytest.mark.parametrize("n_streams", [WVSPLITKRC_SLOTS, WVSPLITKRC_SLOTS + 4])
265 +@pytest.mark.parametrize("dtype", DTYPES)
266 +@pytest.mark.parametrize("seed", SEEDS)
267 +@pytest.mark.skipif(not current_platform.is_rocm(), reason="only test for rocm")
268 +@pytest.mark.skipif(not on_gfx950(), reason="only meant for gfx950")
269 +def test_rocm_wvsplitkrc_multistream(n_streams, dtype, seed):
270 + """Concurrent streams must not share a split-K workspace.
271 +
272 + The workspace is cleared by protocol rather than per invocation, so sharing
273 + one corrupts the reduction and returns wrong numbers rather than failing.
274 + n_streams > WVSPLITKRC_SLOTS additionally exercises the overflow path.
275 + """
276 + torch.manual_seed(seed)
277 +
278 + n, k, m = 16, 2048, 256
279 + iters = 4
280 + xavier = math.sqrt(2 / k)
281 +
282 + B = torch.randn(m, k, dtype=dtype, device="cuda") * xavier
283 + As = [
284 + torch.randn(n, k, dtype=dtype, device="cuda") * xavier for _ in range(n_streams)
285 + ]
286 + refs = [torch.nn.functional.linear(A, B, None) for A in As]
287 + torch.accelerator.synchronize()
288 +
289 + streams = [torch.Stream() for _ in range(n_streams)]
290 + outs: list[list[torch.Tensor]] = [[] for _ in range(n_streams)]
291 + _release_together(streams)
292 + # Round-robin the enqueues so kernels from different streams are in flight
293 + # at the same time.
294 + for _ in range(iters):
295 + for i, s in enumerate(streams):
296 + with torch.cuda.stream(s):
297 + outs[i].append(ops.wvSplitKrc(As[i], B, WVSPLITKRC_TEST_CU, None))
298 + for s in streams:
299 + torch.accelerator.current_stream().wait_stream(s)
300 + torch.accelerator.synchronize()
301 +
302 + for i in range(n_streams):
303 + for out in outs[i]:
304 + torch.testing.assert_close(out, refs[i], atol=1e-3, rtol=1e-2)
305 +
306 +
307 +def _release_together(streams):
308 + """Release every stream at once, once everything below is enqueued.
309 +
310 + An event recorded on an idle stream is already complete, so waiting on it
311 + gates nothing -- the spin, not the event, is what holds the streams.
312 + """
313 + blocker = torch.Stream()
314 + with torch.cuda.stream(blocker):
315 + torch.cuda._sleep(_GATE_CYCLES)
316 + gate = torch.Event()
317 + gate.record()
318 + for s in streams:
319 + s.wait_event(gate)
320 +
321 +
322 @pytest.mark.parametrize("n,k,m", NKM_FACTORS_LLMM1)
323 @pytest.mark.parametrize("dtype", DTYPES)
324 @pytest.mark.parametrize("rows_per_block", [2, 4, 8, 16])
vllm/distributed/parallel_state.py
+6
@@ -2357,6 +2357,12 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):
2357 from vllm.platforms import current_platform
2358
2359 if not current_platform.is_cpu():
2360 + from vllm.triton_utils import HAS_TRITON
2361 +
2362 + if HAS_TRITON:
2363 + from vllm.v1.sample.ops.topk_topp_triton import reset_buffer_cache
2364 +
2365 + reset_buffer_cache()
2366 torch.accelerator.empty_cache()
2367 try:
2368 torch.accelerator.empty_host_cache()
vllm/model_executor/layers/fused_moe/runner/moe_runner.py
-13
@@ -283,24 +283,11 @@ class MoERunner(MoERunnerInterface):
283 self._shared_experts: SharedExperts | None = None
284 if shared_experts is not None:
285 can_overlap = lambda: self._quant_method.mk_can_overlap_shared_experts
286 - # When unquantized, shared expert inputs alias the hidden states,
287 - # which can lead to race condition in multi-stream mode. Quantized
288 - # routed experts copy the input into a fresh buffer first, breaking
289 - # the alias, so overlap is safe. Only observed on ROCm.
290 - routed_input_is_quantized = lambda: (
291 - self.routed_experts.quant_method.moe_quant_config is not None
292 - and self.routed_experts.quant_method.moe_quant_config.quant_dtype
293 - is not None
294 - )
295 - is_multistream_safe = lambda: (
296 - not current_platform.is_rocm() or routed_input_is_quantized()
297 - )
286 self._shared_experts = SharedExperts(
287 shared_experts,
288 moe_config=moe_config,
289 enable_dbo=enable_dbo,
290 mk_can_overlap_shared_experts=can_overlap,
303 - is_multistream_safe=is_multistream_safe,
291 )
292
293 # Needed for string -> MoERunner layer lookup in custom ops.
vllm/model_executor/layers/fused_moe/runner/shared_experts.py
-13
@@ -43,7 +43,6 @@ class SharedExperts(torch.nn.Module):
43 moe_config: FusedMoEConfig,
44 enable_dbo: bool,
45 mk_can_overlap_shared_experts: Callable[[], bool],
46 - is_multistream_safe: Callable[[], bool],
46 ):
47 super().__init__()
48
@@ -58,10 +57,6 @@ class SharedExperts(torch.nn.Module):
57
58 self._mk_can_overlap_shared_experts = mk_can_overlap_shared_experts
59
61 - # Might not be safe to run multi-stream mode if routed and shared experts
62 - # alias the same inputs
63 - self._is_multistream_safe = is_multistream_safe
64 -
60 # Allow disabling of the separate shared experts stream for
61 # debug purposes.
62 # TODO: Remove this after more extensive testings with TP/DP
@@ -112,19 +107,11 @@ class SharedExperts(torch.nn.Module):
107 if self._mk_can_overlap_shared_experts():
108 return SharedExpertsOrder.MK_INTERNAL_OVERLAPPED
109
115 - # On ROCm, empirically only DP-only deployments benefit from the overlap.
116 - overlap_is_beneficial = not current_platform.is_rocm() or (
117 - self._moe_config.moe_parallel_config.dp_size > 1
118 - and self._moe_config.moe_parallel_config.tp_size == 1
119 - )
120 -
110 should_run_shared_in_aux_stream = (
111 current_platform.is_cuda_alike()
112 and self._stream is not None
113 and hidden_states.shape[0]
114 <= envs.VLLM_SHARED_EXPERTS_STREAM_TOKEN_THRESHOLD
126 - and overlap_is_beneficial
127 - and self._is_multistream_safe()
115 )
116
117 if should_run_shared_in_aux_stream:
vllm/model_executor/layers/utils.py
+50 -19
@@ -249,6 +249,45 @@ def use_aiter_triton_gemm(n, m, k, dtype):
249 )
250
251
252 +def wvsplitkrc_dispatch(n: int, k: int, m: int, cu_count: int) -> tuple[int, bool]:
253 + """Pick the K-shard split for wvSplitKrc and say whether the shape fits.
254 +
255 + Mirrors wvSplitKrc() in csrc/rocm/skinny_gemms.cu, which is also where the
256 + shard cap is explained. Both must pick the same chunkk or the workspace
257 + check here bounds the wrong k_rnd.
258 +
259 + Returns:
260 + The CHUNKK the kernel will dispatch with, and whether the CU budget and
261 + split-K workspace admit the shape at all.
262 + """
263 + # Next ^2 of n
264 + N_p2 = 1 << (n - 1).bit_length()
265 + # How many of 4 waves in a group can work on same 16 Ms at same time?
266 + # This reduces the Ms each group works on, i.e. increasing the CUs needed.
267 + GrpsShrB = min(N_p2 // 16, 4)
268 + # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile), and each
269 + # working on a 512-shard of K, how many CUs would we need?
270 + CuNeeded = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512) * GrpsShrB
271 +
272 + CHUNKK2_MAX_SHARDS = 11
273 + shards_chunkk2 = (k + 256 - 1) // 256 # 256-wide shards
274 + chunkk = (
275 + 2
276 + if (
277 + N_p2 != 16
278 + and CuNeeded * 2 <= cu_count
279 + and shards_chunkk2 <= CHUNKK2_MAX_SHARDS
280 + )
281 + else 1
282 + )
283 +
284 + # Deterministic reduction stores one fp32 partial per (M, N, k-shard); all
285 + # of them must fit the split-K workspace.
286 + k_rnd = (k + 512 // chunkk - 1) // (512 // chunkk)
287 + fits = N_p2 * m * k_rnd <= 128 * 1024 * 12 and CuNeeded <= cu_count
288 + return chunkk, fits
289 +
290 +
291 def rocm_unquantized_gemm_impl(
292 x: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor | None = None
293 ) -> torch.Tensor:
@@ -260,21 +299,7 @@ def rocm_unquantized_gemm_impl(
299
300 cu_count = num_compute_units()
301
263 - # Next ^2 of n
264 - N_p2 = 1 << (n - 1).bit_length()
265 - # With 64 Ms per CU (each of 4 SIMDs working on a 16x16 tile),
266 - # and each working on a 512-shard of K, how many CUs would we need?
267 - rndup_cus = ((m + 64 - 1) // 64) * ((k + 512 - 1) // 512)
268 - # How many of 4 waves in a group can work on same 16 Ms at same time?
269 - # This reduces the Ms each group works on, i.e. increasing the number of CUs needed.
270 - GrpsShrB = min(N_p2 // 16, 4)
271 - # Given the above, how many CUs would we need?
272 - CuNeeded = rndup_cus * GrpsShrB
273 - # Deterministic reduction stores one float workspace value per K shard.
274 - fits_wvsplitkrc = (
275 - N_p2 * m * ((k + 512 - 1) // 512)
276 - ) <= 128 * 1024 * 12 # deterministic
277 - fits_wvsplitkrc &= CuNeeded <= cu_count
302 + _, fits_wvsplitkrc = wvsplitkrc_dispatch(n, k, m, cu_count)
303
304 skinny_operands_compatible = weight.is_contiguous() and (
305 bias is None or bias.is_contiguous()
@@ -362,12 +387,18 @@ direct_register_custom_op(
387
388 @functools.cache
389 def warmup_rocm_skinny_gemm_workspaces(device: torch.device) -> None:
365 - """Eagerly allocate wvSplitKrc's process-lifetime static workspaces.
390 + """Eagerly allocate wvSplitKrc's per-device split-K workspace pool.
391 +
392 + wvSplitKrc partitions one per-device allocation into ``kWvSlots`` slots
393 + (csrc/rocm/skinny_gemms.cu) and hands each stream one on first use, so that
394 + two streams never share the split-K partials and counters.
395
367 - They are otherwise created lazily on the first qualifying GEMM
396 + The pool is otherwise created lazily on the first qualifying GEMM
397 (csrc/rocm/skinny_gemms.cu), which can be the first real request — after
369 - the KV cache backing buffer exists. If one landed in that segment's
370 - rounding tail, it would pin the entire segment at engine shutdown.
398 + the KV cache backing buffer exists. If it landed in that segment's rounding
399 + tail, it would pin the entire segment at engine shutdown; it could also land
400 + inside a cudagraph capture, where it would be taken from the graph's private
401 + pool and its zero-fill would become a replayed graph node.
402 """
403 from vllm.platforms.rocm import on_gfx950
404