From 19af173d3af78533964e380242eb849aef94d16a Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Wed, 21 Oct 2026 15:09:41 -0400 Subject: [PATCH 1/3] [Perf][Spec Decode] DSpark: run the Markov walk in one kernel with int8 W2 on chip (Hopper) Add dspark_markov_quantization="int8". On sm_90 the whole DSpark Markov walk runs in one cooperative launch: W2 is quantized to int8 with one fp32 scale per row, most of it stays in shared memory across all SMs and the rest streams from L2, and the per-step argmax across CTAs goes through one atomic {key, count} pair. Sampling runs in the kernel and writes the bf16 logits it drew from into draft_logits for rejection sampling. Off by default. The kernel is adapted from MIT-licensed int8 Markov walk kernels; the notice is kept in csrc/dspark_markov_walk.cu. Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com> --- .buildkite/test_areas/kernels.yaml | 12 + CMakeLists.txt | 14 + .../kernels/benchmark_dspark_markov_walk.py | 145 +++ csrc/dspark_markov_walk.cu | 901 ++++++++++++++++++ csrc/torch_bindings.cpp | 10 + tests/kernels/test_dspark_markov_walk.py | 235 +++++ .../acceptance_rates/dspark/test_dspark.py | 21 +- vllm/_custom_ops.py | 45 + vllm/config/speculative.py | 17 + .../gpu/spec_decode/dspark/markov_walk.py | 270 ++++++ .../gpu/spec_decode/dspark/speculator.py | 78 ++ 11 files changed, 1747 insertions(+), 1 deletion(-) create mode 100644 benchmarks/kernels/benchmark_dspark_markov_walk.py create mode 100644 csrc/dspark_markov_walk.cu create mode 100644 tests/kernels/test_dspark_markov_walk.py create mode 100644 vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py diff --git a/.buildkite/test_areas/kernels.yaml b/.buildkite/test_areas/kernels.yaml index bbec4403fb74..5798bd763139 100644 --- a/.buildkite/test_areas/kernels.yaml +++ b/.buildkite/test_areas/kernels.yaml @@ -100,6 +100,18 @@ steps: - tests/kernels/core/test_minimax_reduce_rms.py - tests/utils.py +- label: ":nvidia: (H100) DSpark Markov Walk Kernel" + key: dspark-markov-walk-kernel-test-h100 + timeout_in_minutes: 20 + device: h100 + source_file_dependencies: + - csrc/dspark_markov_walk.cu + - csrc/torch_bindings.cpp + - vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py + - tests/kernels/test_dspark_markov_walk.py + commands: + - pytest -v -s kernels/test_dspark_markov_walk.py + - label: ":nvidia: (H100) DeepSeek V4 Kernels" key: deepseek-v4-kernel-test-h100 timeout_in_minutes: 30 diff --git a/CMakeLists.txt b/CMakeLists.txt index 9f586f4b0483..228253d655e0 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -787,6 +787,20 @@ if(VLLM_GPU_LANG STREQUAL "CUDA" OR VLLM_GPU_LANG STREQUAL "HIP") "(requires SM90+ and CUDA >= 12.0).") endif() + # DSpark Markov walk with int8 W2 kept on chip (wgmma, sm_90a only) + cuda_archs_loose_intersection(DSPARK_MARKOV_WALK_ARCHS "9.0a" "${CUDA_ARCHS}") + if(${CMAKE_CUDA_COMPILER_VERSION} VERSION_GREATER_EQUAL 12.0 AND DSPARK_MARKOV_WALK_ARCHS) + set(DSPARK_MARKOV_WALK_SRCS "csrc/dspark_markov_walk.cu") + set_gencode_flags_for_srcs( + SRCS "${DSPARK_MARKOV_WALK_SRCS}" + CUDA_ARCHS "${DSPARK_MARKOV_WALK_ARCHS}") + list(APPEND VLLM_STABLE_EXT_SRC "${DSPARK_MARKOV_WALK_SRCS}") + message(STATUS "Building dspark_markov_walk for archs: ${DSPARK_MARKOV_WALK_ARCHS}") + else() + message(STATUS "Not building dspark_markov_walk as no compatible archs found " + "(requires sm_90a and CUDA >= 12.0).") + endif() + # # CUTLASS scaled_mm kernels (moved from _C to _C_stable_libtorch) # diff --git a/benchmarks/kernels/benchmark_dspark_markov_walk.py b/benchmarks/kernels/benchmark_dspark_markov_walk.py new file mode 100644 index 000000000000..ee0bd844b461 --- /dev/null +++ b/benchmarks/kernels/benchmark_dspark_markov_walk.py @@ -0,0 +1,145 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""DSpark Markov walk: the eager per-step loop vs dspark_markov_walk (sm_90a). + +Both arms walk the same base logits. Times are per call (all draft steps), +replayed from a CUDA graph with a cold L2 cache. +""" + +from functools import partial +from statistics import median + +import torch +import torch.nn.functional as F + +from vllm.utils.argparse_utils import FlexibleArgumentParser +from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample +from vllm.v1.worker.gpu.spec_decode.dspark.markov_walk import DSparkMarkovWalker + + +def eager_walk( + w1, w2, base, anchor, tokens, idx, temperature, seeds, pos, draft_logits, cols +): + """The loop of DSparkSpeculator._sample_sequential (full vocabulary).""" + prev = anchor + for i in range(base.shape[1]): + logits = base[:, i] + F.linear(F.embedding(prev, w1), w2) + if draft_logits is None: + tok = logits.argmax(dim=-1) + else: + tok = gumbel_sample( + logits, + idx[:, i], + temperature, + seeds, + pos[:, i] - 1, + apply_temperature=True, + is_drafting=True, + logits_cache=draft_logits, + logits_cache_col=cols[i], + ) + tokens[:, i] = tok + prev = tok + + +def timer(method: str): + if method == "cupti": + from flashinfer.testing import bench_gpu_time_with_cupti + + return partial( + bench_gpu_time_with_cupti, use_cuda_graph=True, cold_l2_cache=True + ) + from flashinfer.testing import bench_gpu_time_with_cudagraph + + return partial(bench_gpu_time_with_cudagraph, cold_l2_cache=True) + + +@torch.inference_mode() +def main(args) -> None: + device = torch.device("cuda") + bench = timer(args.timer) + print("vocab steps batch mode | eager_us fused_us speedup | eager W2 GB/s") + for vocab in args.vocab: + g = torch.Generator(device=device).manual_seed(0) + w1 = (torch.randn(vocab, 256, generator=g, device=device) * 0.05).bfloat16() + w2 = (torch.randn(vocab, 256, generator=g, device=device) * 0.02).bfloat16() + for steps in args.steps: + walker = DSparkMarkovWalker(w1, w2, steps) + walker.warmup() + max_reqs = max(args.batch) + draft_logits = torch.zeros( + max_reqs, steps, vocab, dtype=torch.bfloat16, device=device + ) + temperature = torch.ones(max_reqs, device=device) + seeds = torch.arange(max_reqs, dtype=torch.int64, device=device) + cols = torch.arange(steps, dtype=torch.int32, device=device) + for batch in args.batch: + base = ( + torch.randn(batch, steps, vocab, generator=g, device=device) * 3 + ).bfloat16() + anchor = torch.randint(0, vocab, (batch,), generator=g, device=device) + anchor_ids = anchor.int() + idx = ( + torch.arange(batch, device=device)[:, None] + .expand(batch, steps) + .contiguous() + ) + pos = torch.arange(batch * steps, device=device).view(batch, steps) + 1 + tokens = torch.empty(batch, steps, dtype=torch.int64, device=device) + for mode in args.mode: + cache = draft_logits if mode == "probabilistic" else None + temp = temperature if mode == "probabilistic" else None + + run_eager = partial( + eager_walk, + w1, + w2, + base, + anchor, + tokens, + idx, + temperature, + seeds, + pos, + cache, + cols, + ) + run_fused = partial( + walker.walk, + base, + anchor_ids, + 1, + idx.view(-1), + temp, + seeds, + pos.view(-1), + tokens, + cache, + ) + eager_us = median(bench(run_eager)) * 1e3 + fused_us = median(bench(run_fused)) * 1e3 + gbps = steps * vocab * 256 * 2 / (eager_us * 1e3) + print( + f"{vocab} {steps} {batch} {mode} | {eager_us:.1f} " + f"{fused_us:.1f} {eager_us / fused_us:.2f}x | " + f"{gbps:.0f}" + ) + + +if __name__ == "__main__": + parser = FlexibleArgumentParser(description=__doc__) + parser.add_argument( + "--vocab", type=int, nargs="+", default=[151936, 129280, 248320] + ) + parser.add_argument("--steps", type=int, nargs="+", default=[5, 7, 16]) + parser.add_argument( + "--batch", type=int, nargs="+", default=[1, 2, 4, 8, 16, 32, 64] + ) + parser.add_argument( + "--mode", + nargs="+", + default=["greedy", "probabilistic"], + choices=["greedy", "probabilistic"], + ) + parser.add_argument("--timer", default="cupti", choices=["cupti", "cudagraph"]) + main(parser.parse_args()) diff --git a/csrc/dspark_markov_walk.cu b/csrc/dspark_markov_walk.cu new file mode 100644 index 000000000000..b2d23481fc4b --- /dev/null +++ b/csrc/dspark_markov_walk.cu @@ -0,0 +1,901 @@ +// SPDX-License-Identifier: Apache-2.0 +// SPDX-FileCopyrightText: Copyright contributors to the vLLM project +// +// The DSpark Markov walk in one persistent launch, int8 W2 kept on chip +// (Hopper, sm_90a). +// +// Adapted from markov_walk_i8s.cu of Xinwei Qiang's int8 DSpark Markov walk +// kernels (https://github.com/tie-pilot-qxw/dspark-markov-walk), which carry +// this notice: +// +// MIT License +// +// Copyright (c) 2026 tie-pilot-qxw +// +// Permission is hereby granted, free of charge, to any person obtaining a +// copy of this software and associated documentation files (the +// "Software"), to deal in the Software without restriction, including +// without limitation the rights to use, copy, modify, merge, publish, +// distribute, sublicense, and/or sell copies of the Software, and to permit +// persons to whom the Software is furnished to do so, subject to the +// following conditions: +// +// The above copyright notice and this permission notice shall be included +// in all copies or substantial portions of the Software. +// +// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS +// OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF +// MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN +// NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, +// DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR +// OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE +// USE OR OTHER DEALINGS IN THE SOFTWARE. +// +// What it computes, for every request b of the batch: +// +// prev = anchor[b] +// for k in range(num_steps): +// logit = bf16(base[b, k] + W2 @ W1[prev]) # [V] +// tok = argmax(logit) or a sample of softmax(logit / T[b]) +// prev = tok +// +// The eager loop launches a full-vocabulary GEMV, an add and a sampler per +// step and reads all of W2 from HBM every step. Here one cooperative grid +// (one CTA per SM) runs every step: +// +// * W2 is int8 with one fp32 scale per row. Each CTA owns kTiles 64-row tiles: +// kRes of them stay in SMEM for the whole call, the rest stream from L2 +// through a kRing-slot TMA ring every step. W1[prev] is stored as two int8 +// planes, u ~= s_hi * q_hi + s_lo * q_lo with s_lo = s_hi / 256, so both dot +// products are exact int32 sums on int8 tensor cores (wgmma m64n64k32, u as +// the register A operand, W2 as the SMEM B operand). Nothing is dequantized. +// * The logits are rounded to bf16 once. Argmax and sampler read those values +// and the same bits go to the draft-logit cache, so the rejection sampler +// sees exactly the distribution a sampled token was drawn from. +// * Each step ends with one global argmax per request: every CTA folds its +// best (order-preserving fp32 value, inverted row) key into a {key, count} +// pair with atomicMax and a release add, then polls the pair. Ties go to +// the smallest row, as torch.argmax. The pairs of a call alternate between +// two sets by call parity, so nothing needs a memset between CUDA graph +// replays. +// * Sampling is exact up to floating-point error: inside each 8-row item the +// row is drawn by inverse CDF, across items by Gumbel-max on the item's +// log2 mass. The noise is Philox keyed by the request's seed and counted by +// (item, step, position), so a seeded request replays its draws. +// +// Contract with the caller (vllm/v1/worker/gpu/spec_decode/dspark/): +// * base: bf16 [nb, kb >= num_steps, ld] contiguous, ld % 8 == 0. Rows past +// valid_rows (the last CTA's padding) are never read or written. +// * anchor of request b: anchor_ids[b * anchor_stride]. +// * idx_mapping[b * num_steps]: the request-state index of request b, < 0 for +// a padding row. temperature / seeds are indexed by it. temperature is +// optional: without it every request is greedy. +// * pos[b * num_steps + k]: the position that keys the draw of (b, k). +// * tokens[b * tokens_stride + k] receives the walk. +// * draft_logits (optional): bf16 [num_reqs_max, cols, V], unit stride in the +// last dim. Row (idx, k) receives the logits of every sampling request +// (temperature > 0). Greedy requests leave their rows untouched. +// * state: int64 [kStateU64], 128-B aligned, zeroed once and kept. + +#include +#include +#include + +#include "core/registration.h" +#include "torch_utils.h" + +#include +#include + +#include +#include + +namespace vllm { +namespace dspark_markov_walk { + +using u64 = unsigned long long; + +constexpr int kMaxSteps = 16; +constexpr int kMaxReqs = 64; +constexpr int kThreads = 512; +constexpr int kWarpGroups = 4; +constexpr int kBlockReqs = 32; // requests per warpgroup pass (wgmma M = 64) +constexpr int kTileRows = 64; // W2 rows per tile (wgmma N) +constexpr int kTileBytes = kTileRows * 256; +constexpr int kChunkBytes = 2048; // one k32 chunk of a tile +constexpr int kW1RowBytes = 528; // fragment-ordered q_hi | q_lo, s_hi, s_lo +constexpr int kRes = 11; // SMEM-resident tiles per CTA +constexpr int kRing = 2; // TMA ring slots for the streamed tiles +constexpr int kPairU64 = 16; // one {key, count} pair per 128-B line +constexpr int kSetU64 = kMaxSteps * kMaxReqs * kPairU64; +constexpr int kStateU64 = 16 + 2 * kSetU64; +constexpr unsigned kNoiseTag = 0x4D57A1Bu; + +// ---- memory-model helpers +__device__ __forceinline__ u64 ld_relaxed(const u64* p) { + u64 v; + asm volatile("ld.relaxed.gpu.global.u64 %0, [%1];" + : "=l"(v) + : "l"(p) + : "memory"); + return v; +} +// {key, count} in one 16-B load +__device__ __forceinline__ ulonglong2 ld_relaxed_v2(const u64* p) { + ulonglong2 v; + asm volatile("ld.relaxed.gpu.global.v2.u64 {%0, %1}, [%2];" + : "=l"(v.x), "=l"(v.y) + : "l"(p) + : "memory"); + return v; +} +__device__ __forceinline__ void st_relaxed(u64* p, u64 v) { + asm volatile("st.relaxed.gpu.global.u64 [%0], %1;" ::"l"(p), "l"(v) + : "memory"); +} +__device__ __forceinline__ void red_add_release(u64* p, u64 v) { + asm volatile("red.release.gpu.global.add.u64 [%0], %1;" ::"l"(p), "l"(v) + : "memory"); +} + +// order-preserving fp32 in the high word, inverted row in the low word: +// the max key is the best value, ties go to the smallest row +__device__ __forceinline__ u64 pack_key(float v, unsigned row) { + unsigned b = __float_as_uint(v); + b = (b & 0x80000000u) ? ~b : (b | 0x80000000u); + return (static_cast(b) << 32) | static_cast(0xFFFFFFFFu - row); +} +__device__ __forceinline__ unsigned key_row(u64 key) { + return 0xFFFFFFFFu - static_cast(key & 0xFFFFFFFFull); +} + +// Philox4x32-10 +__device__ __forceinline__ uint4 philox(uint4 c, uint2 k) { +#pragma unroll + for (int i = 0; i < 10; ++i) { + const unsigned hi0 = __umulhi(0xD2511F53u, c.x), lo0 = 0xD2511F53u * c.x; + const unsigned hi1 = __umulhi(0xCD9E8D57u, c.z), lo1 = 0xCD9E8D57u * c.z; + c = make_uint4(hi1 ^ c.y ^ k.x, lo1, hi0 ^ c.w ^ k.y, lo0); + k.x += 0x9E3779B9u; + k.y += 0xBB67AE85u; + } + return c; +} + +// lg2.approx without the denormal fix-up: every argument below is normal, 0 +// or inf, where the ftz variant returns the same bits +__device__ __forceinline__ float lg2_ftz(float x) { + float y; + asm("lg2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +__device__ __forceinline__ float ex2_ftz(float x) { + float y; + asm("ex2.approx.ftz.f32 %0, %1;" : "=f"(y) : "f"(x)); + return y; +} +// base-2 Gumbel up to a constant: -log2(-log2(1 - v)), small v by series +__device__ __forceinline__ float gumbel2(unsigned x) { + const float v = (__uint2float_rn(x) + 0.5f) * 2.3283064365386963e-10f; + const float e = v < 0.0078125f + ? v * 1.4426950408889634f * + (1.f + v * (0.5f + v * (0.33333334f + v * 0.25f))) + : -lg2_ftz(1.f - v); + return -lg2_ftz(e); +} + +// An 8-row run of bf16 logits as 4 packed pairs (pair i = rows 2i, 2i + 1): +// its max in both halves, and the first row holding a value +__device__ __forceinline__ unsigned bf16x8_max(unsigned p0, unsigned p1, + unsigned p2, unsigned p3) { + unsigned m; + asm("{\n\t.reg .b32 a, b;\n\t" + "max.bf16x2 a, %1, %2;\n\tmax.bf16x2 b, %3, %4;\n\tmax.bf16x2 a, a, " + "b;\n\t" + "prmt.b32 b, a, a, 0x1032;\n\tmax.bf16x2 %0, a, b;\n\t}" + : "=r"(m) + : "r"(p0), "r"(p1), "r"(p2), "r"(p3)); + return m; +} +__device__ __forceinline__ int bf16x8_first_eq(unsigned p0, unsigned p1, + unsigned p2, unsigned p3, + unsigned m) { + int qi; + asm("{\n\t.reg .pred e0, e1, e2, e3, e4, e5, e6, e7;\n\t" + "setp.eq.bf16x2 e0|e1, %1, %5;\n\tsetp.eq.bf16x2 e2|e3, %2, %5;\n\t" + "setp.eq.bf16x2 e4|e5, %3, %5;\n\tsetp.eq.bf16x2 e6|e7, %4, %5;\n\t" + "mov.b32 %0, 7;\n\t@e6 mov.b32 %0, 6;\n\t@e5 mov.b32 %0, 5;\n\t" + "@e4 mov.b32 %0, 4;\n\t@e3 mov.b32 %0, 3;\n\t@e2 mov.b32 %0, 2;\n\t" + "@e1 mov.b32 %0, 1;\n\t@e0 mov.b32 %0, 0;\n\t}" + : "=r"(qi) + : "r"(p0), "r"(p1), "r"(p2), "r"(p3), "r"(m)); + return qi; +} + +// ---- loads and stores +__device__ __forceinline__ uint4 ldg_nc(const void* p) { + uint4 v; + asm volatile("ld.global.nc.v4.u32 {%0, %1, %2, %3}, [%4];" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p)); + return v; +} +__device__ __forceinline__ uint4 ldg_hint(const void* p, u64 pol) { + uint4 v; + asm volatile("ld.global.L2::cache_hint.v4.u32 {%0, %1, %2, %3}, [%4], %5;" + : "=r"(v.x), "=r"(v.y), "=r"(v.z), "=r"(v.w) + : "l"(p), "l"(pol)); + return v; +} +__device__ __forceinline__ void stg_hint(void* p, uint4 v, u64 pol) { + asm volatile( + "st.global.L2::cache_hint.v4.u32 [%0], {%1, %2, %3, %4}, %5;" ::"l"(p), + "r"(v.x), "r"(v.y), "r"(v.z), "r"(v.w), "l"(pol) + : "memory"); +} +__device__ __forceinline__ u64 evict_first_policy() { + u64 pol; + asm volatile("createpolicy.fractional.L2::evict_first.b64 %0, 1.0;" + : "=l"(pol)); + return pol; +} +__device__ __forceinline__ u64 evict_last_policy() { + u64 pol; + asm volatile("createpolicy.fractional.L2::evict_last.b64 %0, 1.0;" + : "=l"(pol)); + return pol; +} + +// ---- mbarrier + TMA bulk copies +__device__ __forceinline__ unsigned smem_u32(const void* p) { + return static_cast(__cvta_generic_to_shared(p)); +} +__device__ __forceinline__ void mbar_init(unsigned addr, unsigned count) { + asm volatile("mbarrier.init.shared::cta.b64 [%0], %1;" ::"r"(addr), "r"(count) + : "memory"); +} +__device__ __forceinline__ void mbar_expect_tx(unsigned addr, unsigned bytes) { + asm volatile( + "mbarrier.arrive.expect_tx.shared::cta.b64 _, [%0], %1;" ::"r"(addr), + "r"(bytes) + : "memory"); +} +// CTA-scope acquire: only this CTA's TMA wrote the SMEM (a cluster-scope wait +// would also invalidate L1 on every try) +__device__ __forceinline__ void mbar_wait(unsigned addr, unsigned parity) { + unsigned ok = 0; + do { + asm volatile( + "{\n\t.reg .pred p;\n\t" + "mbarrier.try_wait.parity.acquire.cta.shared::cta.b64 p, [%1], %2;\n\t" + "selp.u32 %0, 1, 0, p;\n\t}" + : "=r"(ok) + : "r"(addr), "r"(parity) + : "memory"); + } while (!ok); +} +__device__ __forceinline__ void bulk_g2s(unsigned dst, const void* src, + unsigned bytes, unsigned mbar, + u64 pol) { + asm volatile( + "cp.async.bulk.shared::cluster.global.mbarrier::complete_tx::bytes." + "L2::cache_hint [%0], [%1], %2, [%3], %4;" ::"r"(dst), + "l"(src), "r"(bytes), "r"(mbar), "l"(pol) + : "memory"); +} +__device__ __forceinline__ int ld_volatile_shared(const int* p) { + int v; + asm volatile("ld.volatile.shared.s32 %0, [%1];" + : "=r"(v) + : "r"(smem_u32(p)) + : "memory"); + return v; +} + +// ---- wgmma m64n64k32 s8 x s8 -> s32, A from registers +__device__ __forceinline__ void wg_fence() { + asm volatile("wgmma.fence.sync.aligned;" ::: "memory"); +} +__device__ __forceinline__ void wg_commit() { + asm volatile("wgmma.commit_group.sync.aligned;" ::: "memory"); +} +__device__ __forceinline__ void wg_wait0() { + asm volatile("wgmma.wait_group.sync.aligned 0;" ::: "memory"); +} +// K-major, no swizzle, LBO 128 B, SBO 256 B; the high word is the constant SBO +constexpr unsigned kDescHi = 256 >> 4; +__device__ __forceinline__ unsigned wg_desc_lo(unsigned addr) { + return ((addr & 0x3FFFF) >> 4) | ((128 >> 4) << 16); +} +#define DSPARK_WG_D32 \ + "{%0, %1, %2, %3, %4, %5, %6, %7, %8, %9, %10, %11, %12, %13, %14, %15, " \ + "%16, %17, %18, %19, %20, %21, %22, %23, %24, %25, %26, %27, %28, %29, " \ + "%30, %31}" +#define DSPARK_WG_D32_OUT(d) \ + "+r"(d[0]), "+r"(d[1]), "+r"(d[2]), "+r"(d[3]), "+r"(d[4]), "+r"(d[5]), \ + "+r"(d[6]), "+r"(d[7]), "+r"(d[8]), "+r"(d[9]), "+r"(d[10]), \ + "+r"(d[11]), "+r"(d[12]), "+r"(d[13]), "+r"(d[14]), "+r"(d[15]), \ + "+r"(d[16]), "+r"(d[17]), "+r"(d[18]), "+r"(d[19]), "+r"(d[20]), \ + "+r"(d[21]), "+r"(d[22]), "+r"(d[23]), "+r"(d[24]), "+r"(d[25]), \ + "+r"(d[26]), "+r"(d[27]), "+r"(d[28]), "+r"(d[29]), "+r"(d[30]), \ + "+r"(d[31]) +// d (+)= A B; acc == 0 overwrites +__device__ __forceinline__ void wgmma_rs(int (&d)[32], const uint4& a, + unsigned blo, int acc) { + asm volatile( + "{\n.reg .pred p;\n.reg .b64 bd;\nsetp.ne.b32 p, %37, 0;\n" + "mov.b64 bd, {%36, %38};\n" + "wgmma.mma_async.sync.aligned.m64n64k32.s32.s8.s8 " DSPARK_WG_D32 + ", {%32, %33, %34, %35}, bd, p;\n}" + : DSPARK_WG_D32_OUT(d) + : "r"(a.x), "r"(a.y), "r"(a.z), "r"(a.w), "r"(blo), "r"(acc), + "r"(kDescHi)); +} + +// Layout of one (32-request block, 64-row tile) product: M row 16 w + g is +// q_hi of request 8 w + g of the block, row 16 w + g + 8 its q_lo (the host +// stores W1 rows in that fragment order, so the gather is 8 x LDG.128 into the +// A registers). The host permutes every tile's rows so that accumulator column +// n = 8 i + 2 tig + e is tile row 8 tig + 2 i + e (i < 4) or 32 + 8 tig + +// 2 (i - 4) + e: a thread owns one request and two aligned 8-row runs. +// +// Ring: the streamed tiles of a call form one FIFO q = step * kStr + rank of +// the tile among the streamed ones; each warpgroup consumes its tiles in +// increasing q. The warp whose arrival completes the consumption of q refills +// the slot with q + kRing: it writes the slot's tag first, then arms the +// barrier and issues the copy, and a consumer waits for the tag before the +// parity, so a warpgroup never mistakes fill f - 2 for fill f. +template +__global__ void __launch_bounds__(kThreads, 1) markov_walk_kernel( + const uint8_t* __restrict__ w2_res, const uint8_t* __restrict__ w2_str, + const float* __restrict__ row_scale, const uint8_t* __restrict__ w1f, + const __nv_bfloat16* __restrict__ base, + const int32_t* __restrict__ anchor_ids, int anchor_stride, + const int64_t* __restrict__ idx_mapping, + const float* __restrict__ temperature, const int64_t* __restrict__ seeds, + const int64_t* __restrict__ pos, int64_t* __restrict__ tokens, + int tokens_stride, __nv_bfloat16* __restrict__ draft_logits, + int64_t dl_stride0, int64_t dl_stride1, u64* __restrict__ state, int nb, + int num_steps, int kb, int ld, int valid_rows, unsigned stream_mask) { + constexpr int kRowsCta = kTiles * kTileRows; + constexpr int kStr = kTiles - kRes; + static_assert(kRing >= 1 && kRing <= kStr, "ring depth"); + extern __shared__ __align__(128) unsigned char smem_raw[]; + unsigned char* res_s = smem_raw; // [kRes][16 KiB] + unsigned char* ring_s = res_s + kRes * kTileBytes; // [kRing][16 KiB] + float* scale_s = reinterpret_cast(ring_s + kRing * kTileBytes); + __shared__ __align__(8) u64 mbar_res[kRes + 1]; // + 1: the row scales + __shared__ __align__(8) u64 mbar_full[kRing]; + __shared__ int slot_q[kRing]; // FIFO index the slot holds or awaits + __shared__ int slot_cnt[kRing]; // monotone consumer arrivals + __shared__ u64 wbest_s[kMaxReqs][kWarpGroups]; + __shared__ unsigned tok_s[kMaxReqs]; + // per sampling request and step: Philox key and counter words, and the + // draft-logit row (kept here, not in registers, for the tile loop) + __shared__ uint4 noise_s[kMaxReqs]; + __shared__ __nv_bfloat16* dl_s[kMaxReqs]; + const int lane = threadIdx.x & 31, warp = threadIdx.x >> 5; + const int wg = warp >> 2, wq = warp & 3, g = lane >> 2, tig = lane & 3; + const int row_base = blockIdx.x * kRowsCta; + + const u64 round = ld_relaxed(state); + const int set = static_cast(round & 1); + u64* pairs = state + 16 + set * kSetU64; + auto pair = [&](int k, int b) { + return pairs + kPairU64 * (k * kMaxReqs + b); + }; + // clear the set the next call uses (the previous call, which used it, has + // finished: calls are serialized on one stream) + if (blockIdx.x == 0) + for (int i = threadIdx.x; i < kMaxSteps * kMaxReqs * 2; i += kThreads) + st_relaxed( + state + 16 + (set ^ 1) * kSetU64 + (i >> 1) * kPairU64 + (i & 1), 0); + + // ---- work split: nblk 32-request blocks, wpb warpgroups per block + const int nblk = (nb + kBlockReqs - 1) / kBlockReqs; + const int wpb = kWarpGroups / nblk; + const int blk = wg / wpb, phase = wg % wpb; + const int b = blk * kBlockReqs + 8 * wq + g; // this thread's request + const bool live = b < nb; + // Run 0 (tile rows 8 tig .. +8) / run 1 (32 + 8 tig .. +8) of tile t holds + // vocabulary rows iff t < lim0 / lim1 (valid_rows % 8 == 0). Dead requests + // get 0, which also skips their base loads. + const int vloc = valid_rows - row_base; + const int lim0 = + live ? min(kTiles, max(0, vloc - 8 * tig + kTileRows - 1) / kTileRows) + : 0; + const int lim1 = + live + ? min(kTiles, max(0, vloc - 32 - 8 * tig + kTileRows - 1) / kTileRows) + : 0; + const int consumers = 4 * nblk; // ring arrivals per fill (4 warps / WG) + const int nq = kStr * num_steps; // streamed fills per call + auto str_idx = [&](int t) { return __popc(stream_mask & ((1u << t) - 1u)); }; + + // ---- per request: sampling parameters (temperature <= 0: greedy) + auto request_temp = [&](int i) { + const int64_t req = idx_mapping[static_cast(i) * num_steps]; + return temperature != nullptr && req >= 0 ? temperature[req] : 0.f; + }; + auto stage_step = [&](int i, int k) { // request i's noise and row of step k + const int64_t req = idx_mapping[static_cast(i) * num_steps]; + if (request_temp(i) <= 0.f) return; + const u64 s = static_cast(seeds[req]); + const u64 p = + static_cast(pos[static_cast(i) * num_steps + k]); + noise_s[i] = make_uint4( + static_cast(s), static_cast(s >> 32), + static_cast(p), static_cast(p >> 32) ^ kNoiseTag); + dl_s[i] = draft_logits != nullptr + ? draft_logits + req * dl_stride0 + k * dl_stride1 + : nullptr; + }; + float inv_t = 0.f; + if (live) { + const float t = request_temp(b); + inv_t = t > 0.f ? 1.f / t : 0.f; + } + const float t2 = inv_t * 1.4426950408889634f; + + // ---- once per call: row scales and resident tiles by TMA, ring prefill + const u64 pol_last = evict_last_policy(), pol_first = evict_first_policy(); + auto res_mbar = [&](int i) { return smem_u32(&mbar_res[i]); }; + auto issue_fill = [&](int q) { + const int slot = q % kRing; + slot_q[slot] = q; + const unsigned mb = smem_u32(&mbar_full[slot]); + mbar_expect_tx(mb, kTileBytes); + // the last step's copies are demoted for whatever runs next + bulk_g2s(smem_u32(ring_s + slot * kTileBytes), + w2_str + (static_cast(blockIdx.x) * kStr + q % kStr) * + kTileBytes, + kTileBytes, mb, q >= nq - kStr ? pol_first : pol_last); + }; + if (threadIdx.x == 0) { + for (int i = 0; i <= kRes; ++i) mbar_init(res_mbar(i), 1); + for (int s = 0; s < kRing; ++s) { + mbar_init(smem_u32(&mbar_full[s]), 1); + slot_cnt[s] = 0; + } + asm volatile("fence.mbarrier_init.release.cluster;" ::: "memory"); + mbar_expect_tx(res_mbar(kRes), kRowsCta * 4); + bulk_g2s(smem_u32(scale_s), row_scale + row_base, kRowsCta * 4, + res_mbar(kRes), pol_first); + for (int q = 0; q < kRing && q < nq; ++q) issue_fill(q); + for (int i = 0; i < kRes; ++i) { + mbar_expect_tx(res_mbar(i), kTileBytes); + bulk_g2s( + smem_u32(res_s + i * kTileBytes), + w2_res + (static_cast(blockIdx.x) * kRes + i) * kTileBytes, + kTileBytes, res_mbar(i), pol_first); + } + } + for (int i = threadIdx.x; i < nb; i += kThreads) { + const int a = anchor_ids[static_cast(i) * anchor_stride]; + tok_s[i] = static_cast(min(max(a, 0), valid_rows - 1)); + stage_step(i, 0); + } + __syncthreads(); // barriers initialized; the copies are still in flight + mbar_wait(res_mbar(kRes), 0); + + const unsigned res_desc0 = wg_desc_lo(smem_u32(res_s)); + const unsigned ring_desc0 = wg_desc_lo(smem_u32(ring_s)); + // base of this thread's next tile (bf16 x 8, two runs) + uint4 nb0 = make_uint4(0, 0, 0, 0), nb1 = make_uint4(0, 0, 0, 0); + + for (int k = 0; k < num_steps; ++k) { + // u fragments of this thread's request (dead requests: zeros, so the + // warpgroup's MMA stays well defined) + uint4 a[8]; + float s_lo = 0.f; + if (live) { + const uint8_t* src = w1f + static_cast(tok_s[b]) * kW1RowBytes; +#pragma unroll + for (int j = 0; j < 8; ++j) a[j] = ldg_nc(src + j * 64 + tig * 16); + s_lo = __ldg(reinterpret_cast(src + 516)); + } else { +#pragma unroll + for (int j = 0; j < 8; ++j) a[j] = make_uint4(0, 0, 0, 0); + } + // element offset of (b, k, run 0 of tile 0); < 2^31, checked on the host + const int ob = (b * kb + k) * ld + row_base + 8 * tig; + if (k == 0) { // later steps load their first tile under the exchange + if (phase < lim0) + nb0 = ldg_hint(base + ob + phase * kTileRows, pol_first); + if (phase < lim1) + nb1 = ldg_hint(base + ob + phase * kTileRows + 32, pol_first); + } + float cur_v = -INFINITY; + int cur_r = -1; + +#pragma unroll 2 + for (int t = phase; t < kTiles; t += wpb) { + const bool streamed = (stream_mask >> t) & 1u; + unsigned bd; + int q = 0; + if (streamed) { + q = k * kStr + str_idx(t); + const int slot = q % kRing; + while (ld_volatile_shared(&slot_q[slot]) != q) { + } + mbar_wait(smem_u32(&mbar_full[slot]), + static_cast((q / kRing) & 1)); + bd = ring_desc0 + slot * (kTileBytes >> 4); + } else { + const int ri = t - str_idx(t); + // first use of a resident tile waits for its own copy + if (k == 0) mbar_wait(res_mbar(ri), 0); + bd = res_desc0 + ri * (kTileBytes >> 4); + } + int d[32]; + wg_fence(); +#pragma unroll + for (int j = 0; j < 8; ++j) + wgmma_rs(d, a[j], bd + j * (kChunkBytes >> 4), j); + wg_commit(); + // this tile's base came one tile ago; the next one's loads under the MMA + const int o = ob + t * kTileRows; + const uint4 bb0 = nb0, bb1 = nb1; + if (t + wpb < lim0) nb0 = ldg_hint(base + o + wpb * kTileRows, pol_first); + if (t + wpb < lim1) + nb1 = ldg_hint(base + o + wpb * kTileRows + 32, pol_first); + wg_wait0(); + // keep every read of d below the wait (d is written asynchronously) + asm volatile("" : DSPARK_WG_D32_OUT(d)::"memory"); + if (streamed && lane == 0) { // this warp is done with the slot + const int slot = q % kRing; + const int done = atomicAdd(&slot_cnt[slot], 1) + 1; + if (done == (q / kRing + 1) * consumers && q + kRing < nq) + issue_fill(q + kRing); + } + if (!live) continue; + // bias = s_lo * (256 d_hi + d_lo) (exact int32), logit = row_scale * + // bias + base, rounded to bf16 once + const float4 sa = + *reinterpret_cast(scale_s + t * kTileRows + 8 * tig); + const float4 sb = *reinterpret_cast( + scale_s + t * kTileRows + 8 * tig + 4); + const float4 sc = *reinterpret_cast( + scale_s + t * kTileRows + 32 + 8 * tig); + const float4 sd = *reinterpret_cast( + scale_s + t * kTileRows + 32 + 8 * tig + 4); + const float sr[16] = {sa.x, sa.y, sa.z, sa.w, sb.x, sb.y, sb.z, sb.w, + sc.x, sc.y, sc.z, sc.w, sd.x, sd.y, sd.z, sd.w}; + const unsigned bw[8] = {bb0.x, bb0.y, bb0.z, bb0.w, + bb1.x, bb1.y, bb1.z, bb1.w}; + float lg[16]; // [run][element]: element 2 (i % 4) + e of run i / 4 +#pragma unroll + for (int i = 0; i < 8; ++i) +#pragma unroll + for (int e = 0; e < 2; ++e) { + const int idx = 8 * (i >> 2) + 2 * (i & 3) + e; + const float bias = + s_lo * static_cast(d[4 * i + e] * 256 + d[4 * i + 2 + e]); + const unsigned w = bw[idx >> 1]; + lg[idx] = fmaf(sr[idx], bias, + __uint_as_float(e ? (w & 0xFFFF0000u) : (w << 16))); + } + unsigned pk[8]; // the 16 bf16 logits, packed in memory order +#pragma unroll + for (int p2 = 0; p2 < 8; ++p2) { + const __nv_bfloat162 h = + __floats2bfloat162_rn(lg[2 * p2], lg[2 * p2 + 1]); + pk[p2] = *reinterpret_cast(&h); + } + const int r = row_base + t * kTileRows + 8 * tig; // run 0 rows r .. +8 + if (row_base + (t + 1) * kTileRows > valid_rows) { + // the last CTA's padding runs: bf16 -inf +#pragma unroll + for (int run = 0; run < 2; ++run) + if (r + 32 * run >= valid_rows) { +#pragma unroll + for (int qq = 0; qq < 4; ++qq) pk[4 * run + qq] = 0xFF80FF80u; + } + } + if (inv_t > 0.f) { + __nv_bfloat16* dl = dl_s[b]; + if (dl != nullptr) { + if (t < lim0) + stg_hint(dl + r, make_uint4(pk[0], pk[1], pk[2], pk[3]), pol_first); + if (t < lim1) + stg_hint(dl + r + 32, make_uint4(pk[4], pk[5], pk[6], pk[7]), + pol_first); + } +#pragma unroll + for (int p2 = 0; p2 < 8; ++p2) { + lg[2 * p2] = __uint_as_float(pk[p2] << 16); + lg[2 * p2 + 1] = __uint_as_float(pk[p2] & 0xFFFF0000u); + } + // run 0 and run 1 draw from one Philox block: (x, y) and (z, w) + const uint4 nz = noise_s[b]; + const uint4 xr = + philox(make_uint4(static_cast(r >> 3), + static_cast(k), nz.z, nz.w), + make_uint2(nz.x, nz.y)); +#pragma unroll + for (int run = 0; run < 2; ++run) { + const float* l8 = lg + 8 * run; + const float m = + __uint_as_float(bf16x8_max(pk[4 * run], pk[4 * run + 1], + pk[4 * run + 2], pk[4 * run + 3]) + << 16); + const float mz = m * t2; + float cs[8], c = 0.f; +#pragma unroll + for (int qq = 0; qq < 8; ++qq) { + c += ex2_ftz(fmaf(l8[qq], t2, -mz)); + cs[qq] = c; + } + const unsigned xu = run == 0 ? xr.x : xr.z; + const unsigned xg = run == 0 ? xr.y : xr.w; + const float u = + (static_cast(xu >> 8) + 0.5f) * 5.9604644775390625e-8f; + const float tt = u * c; // < c + int qi = 7; +#pragma unroll + for (int qq = 6; qq >= 0; --qq) qi = tt < cs[qq] ? qq : qi; + const float key_v = mz + lg2_ftz(c) + gumbel2(xg); + if (key_v > cur_v) { + cur_v = key_v; + cur_r = r + 32 * run + qi; + } + } + } else { // greedy on the packed values +#pragma unroll + for (int run = 0; run < 2; ++run) { + const unsigned* p8 = pk + 4 * run; + const unsigned m2 = bf16x8_max(p8[0], p8[1], p8[2], p8[3]); + const int qi = bf16x8_first_eq(p8[0], p8[1], p8[2], p8[3], m2); + const float m = __uint_as_float(m2 << 16); + if (m > cur_v) { + cur_v = m; + cur_r = r + 32 * run + qi; + } + } + } + } + // the 4 tig lanes of a request -> this warpgroup's best for it + { + u64 key = + cur_r >= 0 ? pack_key(cur_v, static_cast(cur_r)) : 0ull; + key = max(key, __shfl_xor_sync(0xffffffffu, key, 1)); + key = max(key, __shfl_xor_sync(0xffffffffu, key, 2)); + if (tig == 0 && live) wbest_s[b][phase] = key; + } + // the next step's first base tile does not depend on the tokens + if (k + 1 < num_steps) { + const int o1 = ob + ld + phase * kTileRows; + if (phase < lim0) nb0 = ldg_hint(base + o1, pol_first); + if (phase < lim1) nb1 = ldg_hint(base + o1 + 32, pol_first); + } + __syncthreads(); + if (warp == 0) { + // lane holds requests lane and lane + 32 +#pragma unroll + for (int qd = 0; qd < 2; ++qd) { + const int bq = lane + 32 * qd; + if (bq < nb) { + u64 bst = 0ull; + for (int p = 0; p < wpb; ++p) bst = max(bst, wbest_s[bq][p]); + atomicMax(pair(k, bq), bst); + red_add_release(pair(k, bq) + 1, 1); + } + } + // both requests of a lane polled in the same round trip + const u64 n = gridDim.x; + ulonglong2 v0 = make_ulonglong2(0, lane < nb ? 0 : n); + ulonglong2 v1 = make_ulonglong2(0, lane + 32 < nb ? 0 : n); + do { + if (v0.y < n) v0 = ld_relaxed_v2(pair(k, lane)); + if (v1.y < n) v1 = ld_relaxed_v2(pair(k, lane + 32)); + } while (v0.y < n || v1.y < n); + const u64 key[2] = {v0.x, v1.x}; +#pragma unroll + for (int qd = 0; qd < 2; ++qd) { + const int bq = lane + 32 * qd; + if (bq < nb) { + // key 0: no row was a candidate anywhere (every logit NaN or -inf, + // e.g. a padded CUDA graph row); emit row 0, as torch.argmax + const unsigned tok = key[qd] != 0ull ? key_row(key[qd]) : 0u; + tok_s[bq] = tok; + if (blockIdx.x == 0) + tokens[static_cast(bq) * tokens_stride + k] = tok; + if (k + 1 < num_steps) { // the next gather's W1 row into L2 + stage_step(bq, k + 1); + const uint8_t* src = w1f + static_cast(tok) * kW1RowBytes; +#pragma unroll + for (int l = 0; l < 4; ++l) + asm volatile( + "prefetch.global.L2::evict_last [%0];" ::"l"(src + 128 * l)); + } + } + } + } + __syncthreads(); + } + if (blockIdx.x == 0 && threadIdx.x == 0) st_relaxed(state, round + 1); +} + +template +void launch(const torch::stable::Tensor& w2_res, + const torch::stable::Tensor& w2_str, + const torch::stable::Tensor& row_scale, + const torch::stable::Tensor& w1f, const torch::stable::Tensor& base, + const int32_t* anchor_ids, int anchor_stride, + const int64_t* idx_mapping, const float* temperature, + const int64_t* seeds, const int64_t* pos, int64_t* tokens, + int tokens_stride, __nv_bfloat16* draft_logits, int64_t dl_stride0, + int64_t dl_stride1, u64* state, int nb, int num_steps, + int valid_rows, unsigned stream_mask, cudaStream_t stream) { + constexpr int kRowsCta = kTiles * kTileRows; + constexpr int kStr = kTiles - kRes; + const int64_t rows_pad = row_scale.numel(); + STD_TORCH_CHECK(rows_pad % kRowsCta == 0, + "row_scale size is not a multiple of the CTA rows"); + const int grid = static_cast(rows_pad / kRowsCta); + STD_TORCH_CHECK(grid <= get_device_prop()->multiProcessorCount, + "dspark_markov_walk: grid of ", grid, + " CTAs exceeds the SM count"); + STD_TORCH_CHECK( + w2_res.numel() == static_cast(grid) * kRes * kTileBytes, + "w2_res size does not match the tiling"); + STD_TORCH_CHECK( + w2_str.numel() == static_cast(grid) * kStr * kTileBytes, + "w2_str size does not match the tiling"); + STD_TORCH_CHECK(__builtin_popcount(stream_mask) == kStr && + (kTiles == 32 || (stream_mask >> kTiles) == 0), + "stream_mask must stream exactly tiles - ", kRes, " tiles"); + const int kb = static_cast(base.size(1)); + const int ld = static_cast(base.size(2)); + // 32-bit element offsets into base, including the dead lanes of a + // partially filled 32-request block and the padding tiles + const int64_t lanes = + static_cast(kBlockReqs) * ((nb + kBlockReqs - 1) / kBlockReqs); + STD_TORCH_CHECK(lanes * kb * ld + rows_pad < (int64_t(1) << 31), + "base is too large for 32-bit offsets"); + const int smem = (kRes + kRing) * kTileBytes + kRowsCta * 4; + auto kernel = markov_walk_kernel; + static bool attr_set = false; + if (!attr_set) { + STD_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); + attr_set = true; + } + cudaLaunchConfig_t cfg = {}; + cfg.gridDim = dim3(grid); + cfg.blockDim = dim3(kThreads); + cfg.dynamicSmemBytes = smem; + cfg.stream = stream; + cudaLaunchAttribute attr[1]; + attr[0].id = cudaLaunchAttributeCooperative; + attr[0].val.cooperative = 1; + cfg.attrs = attr; + cfg.numAttrs = 1; + STD_CUDA_CHECK(cudaLaunchKernelEx( + &cfg, kernel, reinterpret_cast(w2_res.data_ptr()), + reinterpret_cast(w2_str.data_ptr()), + reinterpret_cast(row_scale.data_ptr()), + reinterpret_cast(w1f.data_ptr()), + reinterpret_cast(base.data_ptr()), anchor_ids, + anchor_stride, idx_mapping, temperature, seeds, pos, tokens, + tokens_stride, draft_logits, dl_stride0, dl_stride1, state, nb, num_steps, + kb, ld, valid_rows, stream_mask)); +} + +} // namespace dspark_markov_walk +} // namespace vllm + +// See the header for the layouts; the host side that builds them is +// vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py. +void dspark_markov_walk( + torch::stable::Tensor const& w2_res, torch::stable::Tensor const& w2_str, + torch::stable::Tensor const& row_scale, torch::stable::Tensor const& w1f, + torch::stable::Tensor const& base, torch::stable::Tensor const& anchor_ids, + int64_t anchor_stride, torch::stable::Tensor const& idx_mapping, + std::optional temperature, + torch::stable::Tensor const& seeds, torch::stable::Tensor const& pos, + torch::stable::Tensor& tokens, + std::optional draft_logits, + torch::stable::Tensor& state, int64_t num_steps, int64_t tiles, + int64_t stream_mask) { + namespace dmw = vllm::dspark_markov_walk; + using torch::headeronly::ScalarType; + STD_TORCH_CHECK(base.dim() == 3 && base.is_contiguous() && + base.scalar_type() == ScalarType::BFloat16, + "base must be contiguous bf16 [num_reqs, steps, vocab]"); + const int64_t nb = base.size(0); + const int64_t ld = base.size(2); + STD_TORCH_CHECK(nb >= 1 && nb <= dmw::kMaxReqs, "batch must be in [1, ", + dmw::kMaxReqs, "]"); + STD_TORCH_CHECK(num_steps >= 1 && num_steps <= dmw::kMaxSteps && + base.size(1) >= num_steps, + "num_steps must be in [1, ", dmw::kMaxSteps, "]"); + STD_TORCH_CHECK(ld % 8 == 0, "vocab size must be a multiple of 8"); + const int64_t valid_rows = w1f.size(0); + STD_TORCH_CHECK(valid_rows == ld && valid_rows <= row_scale.numel(), + "W1 rows, W2 rows and the logits' vocab size must match"); + STD_TORCH_CHECK( + w1f.dim() == 2 && w1f.size(1) == dmw::kW1RowBytes && w1f.is_contiguous(), + "w1f must be [vocab, 528] bytes"); + STD_TORCH_CHECK(w2_res.is_contiguous() && w2_str.is_contiguous() && + row_scale.is_contiguous() && + row_scale.scalar_type() == ScalarType::Float, + "bad W2 layout"); + STD_TORCH_CHECK(anchor_ids.scalar_type() == ScalarType::Int && + anchor_ids.numel() >= (nb - 1) * anchor_stride + 1, + "anchor_ids must be int32 with an anchor per request"); + STD_TORCH_CHECK(idx_mapping.scalar_type() == ScalarType::Long && + idx_mapping.is_contiguous() && + idx_mapping.numel() >= nb * num_steps, + "idx_mapping must be int64 [num_reqs * steps]"); + STD_TORCH_CHECK(pos.scalar_type() == ScalarType::Long && + pos.is_contiguous() && pos.numel() >= nb * num_steps, + "pos must be int64 [num_reqs * steps]"); + STD_TORCH_CHECK(seeds.scalar_type() == ScalarType::Long, + "seeds must be int64"); + STD_TORCH_CHECK(tokens.dim() == 2 && tokens.size(0) >= nb && + tokens.size(1) >= num_steps && tokens.stride(1) == 1 && + tokens.scalar_type() == ScalarType::Long, + "tokens must be int64 [>= num_reqs, >= steps]"); + STD_TORCH_CHECK(state.scalar_type() == ScalarType::Long && + state.numel() == dmw::kStateU64 && + reinterpret_cast(state.data_ptr()) % 128 == 0, + "state must be a 128-B aligned int64 [", dmw::kStateU64, "]"); + const float* temp_ptr = nullptr; + if (temperature.has_value()) { + STD_TORCH_CHECK(temperature->scalar_type() == ScalarType::Float, + "temperature must be fp32"); + temp_ptr = reinterpret_cast(temperature->data_ptr()); + } + __nv_bfloat16* dl_ptr = nullptr; + int64_t dl_stride0 = 0, dl_stride1 = 0; + if (draft_logits.has_value()) { + STD_TORCH_CHECK( + draft_logits->dim() == 3 && + draft_logits->scalar_type() == ScalarType::BFloat16 && + draft_logits->stride(2) == 1 && + draft_logits->size(1) >= num_steps && draft_logits->size(2) >= ld && + draft_logits->stride(0) % 8 == 0 && + draft_logits->stride(1) % 8 == 0 && + reinterpret_cast(draft_logits->data_ptr()) % 16 == 0, + "draft_logits must be 16-B aligned bf16 [reqs, cols, " + "vocab] with unit last stride"); + dl_ptr = reinterpret_cast<__nv_bfloat16*>(draft_logits->data_ptr()); + dl_stride0 = draft_logits->stride(0); + dl_stride1 = draft_logits->stride(1); + } + const torch::stable::accelerator::DeviceGuard device_guard( + base.get_device_index()); + auto stream = get_current_cuda_stream(base.get_device_index()); + const unsigned mask = static_cast(stream_mask); + +#define DSPARK_MARKOV_WALK_LAUNCH(TILES) \ + if (tiles == TILES) { \ + dmw::launch( \ + w2_res, w2_str, row_scale, w1f, base, \ + reinterpret_cast(anchor_ids.data_ptr()), \ + static_cast(anchor_stride), \ + reinterpret_cast(idx_mapping.data_ptr()), temp_ptr, \ + reinterpret_cast(seeds.data_ptr()), \ + reinterpret_cast(pos.data_ptr()), \ + reinterpret_cast(tokens.mutable_data_ptr()), \ + static_cast(tokens.stride(0)), dl_ptr, dl_stride0, dl_stride1, \ + reinterpret_cast(state.mutable_data_ptr()), \ + static_cast(nb), static_cast(num_steps), \ + static_cast(valid_rows), mask, stream); \ + return; \ + } + DSPARK_MARKOV_WALK_LAUNCH(16) + DSPARK_MARKOV_WALK_LAUNCH(18) + DSPARK_MARKOV_WALK_LAUNCH(20) + DSPARK_MARKOV_WALK_LAUNCH(24) + DSPARK_MARKOV_WALK_LAUNCH(28) + DSPARK_MARKOV_WALK_LAUNCH(32) +#undef DSPARK_MARKOV_WALK_LAUNCH + STD_TORCH_CHECK(false, "dspark_markov_walk: no build for ", tiles, + " tiles per CTA"); +} + +STABLE_TORCH_LIBRARY_IMPL(_C, CUDA, m) { + m.impl("dspark_markov_walk", TORCH_BOX(&dspark_markov_walk)); +} diff --git a/csrc/torch_bindings.cpp b/csrc/torch_bindings.cpp index f5eba99b86db..202ad4cd910a 100644 --- a/csrc/torch_bindings.cpp +++ b/csrc/torch_bindings.cpp @@ -324,6 +324,16 @@ STABLE_TORCH_LIBRARY_FRAGMENT(_C, ops) { "dsv3_fused_a_gemm(Tensor! output, Tensor mat_a, Tensor mat_b, " "bool enable_pdl=False) -> ()"); + // DSpark Markov walk, int8 W2 on chip (sm_90a); conditionally compiled so + // impl registration is in source file (dspark_markov_walk.cu) + ops.def( + "dspark_markov_walk(Tensor w2_res, Tensor w2_str, Tensor row_scale, " + "Tensor w1_rows, Tensor base_logits, Tensor anchor_ids, " + "int anchor_stride, Tensor idx_mapping, Tensor? temperature, " + "Tensor seeds, Tensor pos, Tensor! draft_tokens, " + "Tensor!? draft_logits, Tensor! state, int num_steps, int tiles, " + "int stream_mask) -> ()"); + // BF16/FP32 x FP32 -> FP32 router GEMM for H=3072, E=256, M<=32 (SM90+). // conditionally compiled so impl registration is in source file ops.def("fp32_router_gemm(Tensor! output, Tensor mat_a, Tensor mat_b) -> ()"); diff --git a/tests/kernels/test_dspark_markov_walk.py b/tests/kernels/test_dspark_markov_walk.py new file mode 100644 index 000000000000..8605077c516d --- /dev/null +++ b/tests/kernels/test_dspark_markov_walk.py @@ -0,0 +1,235 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""dspark_markov_walk against an exact replay of its int8 arithmetic. + +The kernel's dot products are exact int32 sums, so a float64 replay of the +same quantized operands and fp32 roundings reproduces its bf16 logits bit for +bit (up to double rounding in the final fma, which random data never hits). +""" + +import pytest +import torch + +from vllm.platforms import current_platform +from vllm.v1.worker.gpu.spec_decode.dspark.markov_walk import ( + STATE_NUMEL, + DSparkMarkovWalker, + pick_tiles, + quantize_w1, + quantize_w2, + unsupported_reason, +) + +pytestmark = pytest.mark.skipif( + not current_platform.is_cuda() + or not current_platform.is_device_capability(90) + or not hasattr(torch.ops._C, "dspark_markov_walk"), + reason="dspark_markov_walk needs an sm_90 build", +) + +DEVICE = "cuda" +RANK = 256 + + +def make_head(vocab: int, seed: int = 0): + g = torch.Generator(device=DEVICE).manual_seed(seed) + w1 = (torch.randn(vocab, RANK, generator=g, device=DEVICE) * 0.05).bfloat16() + w2 = (torch.randn(vocab, RANK, generator=g, device=DEVICE) * 0.02).bfloat16() + return w1, w2 + + +class Reference: + """Float64 replay of the kernel's numerics.""" + + def __init__(self, w1, w2, logit_scale=1.0): + q_w2, self.row_scale = quantize_w2(w2, logit_scale) + self.q_w2 = q_w2.double() + q_hi, q_lo, _, self.s_lo = quantize_w1(w1) + self.q_hi, self.q_lo = q_hi.double(), q_lo.double() + + def logits(self, base_k: torch.Tensor, prev: torch.Tensor) -> torch.Tensor: + """bf16 [n, V] logits of one step from base [n, V] and prev [n].""" + d_hi = self.q_hi[prev] @ self.q_w2.T + d_lo = self.q_lo[prev] @ self.q_w2.T + combined = (d_hi * 256 + d_lo).long().float().double() + bias = (self.s_lo[prev].double()[:, None] * combined).float().double() + logit = self.row_scale.double()[None] * bias + base_k.double() + return logit.float().bfloat16() + + +def run_walk( + walker, + base, + anchor, + *, + temps=None, + seeds=None, + pos=None, + idx=None, + draft_logits=None, + tokens=None, +): + nb, steps, _ = base.shape + if idx is None: + idx = torch.arange(nb, device=DEVICE).repeat_interleave(steps) + if seeds is None: + seeds = torch.arange(nb, dtype=torch.int64, device=DEVICE) + 1234 + if pos is None: + pos = torch.arange(nb * steps, dtype=torch.int64, device=DEVICE) + if tokens is None: + tokens = torch.full((nb, steps), -1, dtype=torch.int64, device=DEVICE) + walker.walk(base, anchor.int(), 1, idx, temps, seeds, pos, tokens, draft_logits) + return tokens + + +@pytest.mark.parametrize("vocab", [151936, 129280, 248320, 4096]) +@pytest.mark.parametrize("num_reqs", [1, 3, 32, 33, 64]) +@pytest.mark.parametrize("steps", [1, 5, 16]) +def test_greedy_matches_reference(vocab, num_reqs, steps): + if pick_tiles(vocab, current_platform.num_compute_units()) is None: + pytest.skip("vocabulary too large for this GPU") + w1, w2 = make_head(vocab) + walker = DSparkMarkovWalker(w1, w2, steps) + ref = Reference(w1, w2) + g = torch.Generator(device=DEVICE).manual_seed(num_reqs * 100 + steps) + base = ( + torch.randn(num_reqs, steps, vocab, generator=g, device=DEVICE) * 3 + ).bfloat16() + anchor = torch.randint(0, vocab, (num_reqs,), generator=g, device=DEVICE) + tokens = run_walk(walker, base, anchor) + + prev = anchor + for k in range(steps): + expected = ref.logits(base[:, k], prev).argmax(-1) + torch.testing.assert_close(tokens[:, k], expected, rtol=0, atol=0) + prev = tokens[:, k] + + +def test_logit_scale_folds_into_the_bias(): + vocab, steps = 4096, 4 + w1, w2 = make_head(vocab, seed=1) + walker = DSparkMarkovWalker(w1, w2, steps, logit_scale=0.75) + ref = Reference(w1, w2, logit_scale=0.75) + base = (torch.randn(2, steps, vocab, device=DEVICE) * 0.5).bfloat16() + anchor = torch.tensor([3, 4000], device=DEVICE) + tokens = run_walk(walker, base, anchor) + prev = anchor + for k in range(steps): + assert torch.equal(tokens[:, k], ref.logits(base[:, k], prev).argmax(-1)) + prev = tokens[:, k] + + +@pytest.mark.parametrize("num_reqs", [1, 40]) +def test_draft_logits_cache(num_reqs): + """Sampled requests cache the exact logits they drew from; greedy and + padding requests leave their rows alone.""" + vocab, steps, max_reqs = 151936, 7, 128 + w1, w2 = make_head(vocab, seed=2) + walker = DSparkMarkovWalker(w1, w2, steps) + ref = Reference(w1, w2) + base = (torch.randn(num_reqs, steps, vocab, device=DEVICE) * 3).bfloat16() + anchor = torch.randint(0, vocab, (num_reqs,), device=DEVICE) + # request b lives in state slot 2b + 1; every third one is greedy, the + # last one is a padding row + slots = torch.arange(num_reqs, device=DEVICE) * 2 + 1 + slots[-1] = -1 + idx = slots.repeat_interleave(steps) + temperature = torch.zeros(max_reqs, device=DEVICE) + temperature[slots[slots >= 0]] = 0.8 + temperature[slots[::3][slots[::3] >= 0]] = 0.0 + sentinel = 7.0 + cache = torch.full( + (max_reqs, steps, vocab), sentinel, dtype=torch.bfloat16, device=DEVICE + ) + seeds = torch.randint(0, 2**62, (max_reqs,), device=DEVICE) + tokens = run_walk( + walker, + base, + anchor, + temps=temperature, + seeds=seeds, + idx=idx, + draft_logits=cache, + ) + + assert ((tokens >= 0) & (tokens < vocab)).all() + prev = anchor + for k in range(steps): + expected = ref.logits(base[:, k], prev) + for b in range(num_reqs): + slot = int(slots[b]) + sampled = slot >= 0 and float(temperature[slot]) > 0 + if sampled: + assert torch.equal(cache[slot, k], expected[b]) + elif slot >= 0: + assert (cache[slot, k] == sentinel).all() + if not sampled: # greedy and padding rows take the argmax + assert int(tokens[b, k]) == int(expected[b].argmax()) + prev = tokens[:, k] + untouched = torch.ones(max_reqs, dtype=torch.bool, device=DEVICE) + untouched[slots[slots >= 0]] = False + assert (cache[untouched] == sentinel).all() + + +def test_sampling_is_seeded_and_distributed(): + vocab, steps = 4096, 1 + w1 = torch.randn(vocab, RANK, device=DEVICE).bfloat16() + w2 = torch.zeros(vocab, RANK, dtype=torch.bfloat16, device=DEVICE) # no bias + walker = DSparkMarkovWalker(w1, w2, steps) + logits = torch.full((vocab,), -30.0, device=DEVICE) + hot = torch.tensor([5, 77, 1000, 1001, 2047, 4095], device=DEVICE) + logits[hot] = torch.tensor([1.0, 0.5, 0.0, -0.5, 0.25, 0.75], device=DEVICE) + base = logits.bfloat16().view(1, 1, vocab) + temperature = torch.tensor([0.9], device=DEVICE) + probs = torch.softmax(base.float().view(-1) / 0.9, -1)[hot].double() + + n = 6000 + counts = torch.zeros(vocab, dtype=torch.int64, device=DEVICE) + first = None + for i in range(n): + tokens = run_walk( + walker, + base, + torch.zeros(1, device=DEVICE, dtype=torch.long), + temps=temperature, + seeds=torch.tensor([42], device=DEVICE), + pos=torch.tensor([i], device=DEVICE), + ) + counts[tokens.view(-1)] += 1 + if i == 0: + first = tokens.clone() + assert int(counts[hot].sum()) == n + # the same (seed, position) replays the same draw + again = run_walk( + walker, + base, + torch.zeros(1, device=DEVICE, dtype=torch.long), + temps=temperature, + seeds=torch.tensor([42], device=DEVICE), + pos=torch.tensor([0], device=DEVICE), + ) + assert torch.equal(first, again) + expected = probs * n + chi2 = float((((counts[hot].double() - expected) ** 2) / expected).sum()) + assert chi2 < 30.0, (chi2, counts[hot].tolist(), expected.tolist()) # 5 dof + + +def test_non_finite_rows_stay_in_range(): + vocab, steps = 151936, 3 + w1, w2 = make_head(vocab, seed=3) + walker = DSparkMarkovWalker(w1, w2, steps) + base = torch.randn(9, steps, vocab, device=DEVICE).bfloat16() + base[4] = float("nan") + base[7] = float("-inf") + tokens = run_walk(walker, base, torch.randint(0, vocab, (9,), device=DEVICE)) + assert ((tokens >= 0) & (tokens < vocab)).all() + assert (tokens[7] == 0).all() + + +def test_rejects_unsupported_heads(): + w1, w2 = make_head(4096) + assert unsupported_reason(w1, w2, 7) is None + assert unsupported_reason(w1, w2[:, :128], 7) is not None + assert unsupported_reason(w1[:2048], w2, 7) is not None + assert unsupported_reason(w1, w2, 17) is not None + assert STATE_NUMEL == 16 + 2 * 16 * 64 * 16 diff --git a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py index e90100472400..aca60c56cd81 100644 --- a/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py +++ b/tests/v1/e2e/spec_decode/acceptance_rates/dspark/test_dspark.py @@ -1,13 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project -from dataclasses import dataclass +from dataclasses import dataclass, replace import pytest from tests.evals.gsm8k.gsm8k_eval import evaluate_gsm8k_offline from vllm.config import CompilationConfig from vllm.platforms import current_platform +from vllm.v1.worker.gpu.spec_decode.dspark.markov_walk import pick_tiles from ...utils import compute_acceptance_len, compute_acceptance_rate @@ -33,6 +34,7 @@ class DSparkCorrectnessConfig: enforce_eager: bool = False disable_flashinfer_sampler: bool = False language_model_only: bool = False + markov_quantization: str | None = None # References from 12 full GSM8K runs at temperature 1.0: accuracy 0.782-0.814, @@ -46,6 +48,12 @@ class DSparkCorrectnessConfig: attention_backend="FLASH_ATTN", ) +# The one-launch int8 Markov walk keeps the references above; batches stay at +# or below its 64-request limit so every decode step uses it. +QWEN3_DSPARK_DEEPSPEC_INT8_WALK = replace( + QWEN3_DSPARK_DEEPSPEC, max_num_seqs=64, markov_quantization="int8" +) + # References from five 200-question runs at temperature 1.0: accuracy # 0.900-0.955, acceptance rate 0.578-0.595, acceptance length 5.044-5.167. GEMMA4_DSPARK_DEEPSPEC = DSparkCorrectnessConfig( @@ -84,6 +92,15 @@ class DSparkCorrectnessConfig: "config", [ pytest.param(QWEN3_DSPARK_DEEPSPEC, id="qwen3-deepspec"), + pytest.param( + QWEN3_DSPARK_DEEPSPEC_INT8_WALK, + id="qwen3-deepspec-int8-walk", + marks=pytest.mark.skipif( + not current_platform.is_device_capability(90) + or pick_tiles(151936, current_platform.num_compute_units()) is None, + reason="the int8 Markov walk needs a full sm_90 GPU", + ), + ), pytest.param(GEMMA4_DSPARK_DEEPSPEC, id="gemma4-deepspec"), pytest.param(QWEN3_6_DSPARK_SPECULATORS, id="qwen3.6-speculators"), ], @@ -103,6 +120,8 @@ def test_dspark_correctness_and_acceptance_rate( "num_speculative_tokens": config.num_speculative_tokens, "draft_sample_method": "probabilistic", } + if config.markov_quantization is not None: + speculative_config["dspark_markov_quantization"] = config.markov_quantization if config.attention_backend is not None and current_platform.is_cuda(): speculative_config["attention_backend"] = config.attention_backend diff --git a/vllm/_custom_ops.py b/vllm/_custom_ops.py index 586e4d19aa9c..b3f4a5bc4473 100644 --- a/vllm/_custom_ops.py +++ b/vllm/_custom_ops.py @@ -3453,6 +3453,51 @@ def dsv3_fused_a_gemm( torch.ops._C.dsv3_fused_a_gemm(output, mat_a, mat_b, enable_pdl) +def dspark_markov_walk( + w2_res: torch.Tensor, + w2_str: torch.Tensor, + row_scale: torch.Tensor, + w1_rows: torch.Tensor, + base_logits: torch.Tensor, + anchor_ids: torch.Tensor, + anchor_stride: int, + idx_mapping: torch.Tensor, + temperature: torch.Tensor | None, + seeds: torch.Tensor, + pos: torch.Tensor, + draft_tokens: torch.Tensor, + draft_logits: torch.Tensor | None, + state: torch.Tensor, + num_steps: int, + tiles: int, + stream_mask: int, +) -> None: + """DSpark Markov walk in one launch with int8 W2 on chip (sm_90a). + + The layouts come from + ``vllm.v1.worker.gpu.spec_decode.dspark.markov_walk``. + """ + torch.ops._C.dspark_markov_walk( + w2_res, + w2_str, + row_scale, + w1_rows, + base_logits, + anchor_ids, + anchor_stride, + idx_mapping, + temperature, + seeds, + pos, + draft_tokens, + draft_logits, + state, + num_steps, + tiles, + stream_mask, + ) + + def weight_packed_linear_cpu( x: torch.Tensor, packed_weight: torch.Tensor, diff --git a/vllm/config/speculative.py b/vllm/config/speculative.py index 35441ec5b63b..1cf13faae1a0 100644 --- a/vllm/config/speculative.py +++ b/vllm/config/speculative.py @@ -616,6 +616,12 @@ def _resolve_synthetic_acceptance_rates( """For Qwen3 DSpark drafting, evaluate the Markov projection only for the top-k base-logit candidates. Requires draft tensor parallel size 1.""" + dspark_markov_quantization: Literal["int8"] | None = None + """Run DSpark's sequential Markov walk as one persistent kernel that keeps + an int8 copy of the Markov projection on chip, instead of one + full-vocabulary GEMV, add and sample per draft token. sm_90 only. The + target verifies the draft, so only the acceptance length can change.""" + fly_window_size: int | None = Field(default=None, ge=1) """Number of subsequent native acceptance decisions checked by FLy. If unset, defaults to ``min(6, num_speculative_tokens - 1)``.""" @@ -1856,6 +1862,17 @@ def _verify_args(self) -> Self: f"model dtype, got head_dtype={draft.head_dtype}." ) + if self.dspark_markov_quantization is not None: + if self.method != "dspark": + raise ValueError( + "dspark_markov_quantization is only supported by DSpark." + ) + if self.dspark_draft_topk is not None: + raise ValueError( + "dspark_markov_quantization and dspark_draft_topk cannot " + "be combined." + ) + if self.num_speculative_tokens <= 0: raise ValueError( "Expected num_speculative_tokens to be greater " diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py b/vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py new file mode 100644 index 000000000000..3354f665322d --- /dev/null +++ b/vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py @@ -0,0 +1,270 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""One-launch DSpark Markov walk with an int8 Markov projection kept on chip. + +Host side of ``csrc/dspark_markov_walk.cu`` (layouts adapted from Xinwei +Qiang's int8 DSpark Markov walk kernels, MIT License). The eager +walk reads the whole ``V x 256`` projection W2 from HBM once per draft step. +The kernel quantizes W2 to int8 with one fp32 scale per row, keeps most of it +in shared memory across all SMs and streams the rest from L2, and runs all +steps in one cooperative launch. The target model verifies every draft token, +so the only effect of the quantization is on acceptance. +""" + +import torch + +from vllm import _custom_ops as ops +from vllm.logger import init_logger +from vllm.platforms import current_platform + +logger = init_logger(__name__) + +MARKOV_RANK = 256 +MAX_STEPS = 16 +MAX_REQS = 64 +TILE_ROWS = 64 +W1_ROW_BYTES = 528 +STATE_NUMEL = 16 + 2 * MAX_STEPS * MAX_REQS * 16 +# Tiles per CTA built into the kernel; each keeps RES tiles in SMEM and +# streams the others from L2 through a RING-slot ring. +TILE_BUILDS = (16, 18, 20, 24, 28, 32) +RES, RING = 11, 2 +# The timed stream mask of the 18-tile layout (V = 151936 on 132 SMs). +MASK_18 = sum(1 << t for t in (1, 2, 4, 7, 9, 11, 14)) + +# Tile row permutation: accumulator column n of a 64-row tile holds tile row +# rho(n), so every thread owns two aligned 8-row runs. +_RHO = [ + (8 * ((n % 8) // 2) + 2 * (n // 8) + n % 2) + if n // 8 < 4 + else (32 + 8 * ((n % 8) // 2) + 2 * (n // 8 - 4) + n % 2) + for n in range(64) +] + + +def unsupported_reason( + w1: torch.Tensor, w2: torch.Tensor, num_steps: int +) -> str | None: + """Why the kernel cannot walk this Markov head, or None if it can.""" + if not current_platform.is_cuda() or not current_platform.is_device_capability(90): + return "it needs an sm_90 GPU" + if not hasattr(torch.ops._C, "dspark_markov_walk"): + return "this build has no dspark_markov_walk kernel" + if w2.dim() != 2 or w2.shape[1] != MARKOV_RANK: + return f"it needs a rank-{MARKOV_RANK} Markov head" + if w1.shape != w2.shape: + return "it needs a full-vocabulary draft (no draft-to-target map)" + if not w2.dtype.is_floating_point or not w1.dtype.is_floating_point: + return "it needs an unquantized Markov projection" + if w2.shape[0] % 8 != 0: + return "it needs a vocabulary size divisible by 8" + if not 1 <= num_steps <= MAX_STEPS: + return f"it supports at most {MAX_STEPS} draft tokens" + if pick_tiles(w2.shape[0], _num_sms(w2.device)) is None: + return "the vocabulary is too large for this GPU" + return None + + +def _num_sms(device: torch.device) -> int: + return current_platform.num_compute_units(device.index or 0) + + +def pick_tiles(vocab_size: int, num_sms: int) -> int | None: + need = -(-vocab_size // (num_sms * TILE_ROWS)) + return next((t for t in TILE_BUILDS if t >= need), None) + + +def stream_mask(tiles: int) -> int: + """Which tiles stream from L2 (bit t), spread over the four warpgroups. + + With one 32-request block, warpgroup p walks tiles p, p + 4, ...; residents + are shared in proportion to each warpgroup's tile count, the first RING + warpgroups of (1, 2, 3, 0) open on a streamed tile (they consume the ring + prefill) and the others on a resident one. + """ + if tiles == 18: + return MASK_18 + seqs = [list(range(p, tiles, 4)) for p in range(4)] + quota = [RES * len(sq) / tiles for sq in seqs] + res = [int(q) for q in quota] + for _ in range(RES - sum(res)): + parity = [res[0] + res[2], res[1] + res[3]] + p = max( + (p for p in range(4) if res[p] < len(seqs[p])), + key=lambda p: (round(quota[p] - res[p], 9), -parity[p % 2], -p), + ) + res[p] += 1 + starts_streamed = {(1, 2, 3, 0)[i] for i in range(min(RING, 4))} + mask = 0 + for p, sq in enumerate(seqs): + n, nr = len(sq), res[p] + ns = n - nr + if p in starts_streamed: + pos = {i * n // ns for i in range(ns)} if ns else set() + else: + pos = set(range(n)) - ({i * n // nr for i in range(nr)} if nr else set()) + for j in pos: + mask |= 1 << sq[j] + assert bin(mask).count("1") == tiles - RES and mask < (1 << tiles) + return mask + + +def quantize_w2(w2: torch.Tensor, scale: float) -> tuple[torch.Tensor, torch.Tensor]: + """int8 W2 with one fp32 scale per row; ``scale`` (the logit scale) folds in.""" + w = w2.float() + row_scale = (w.abs().amax(1, keepdim=True) / 127.0).clamp_min(1e-30) + q = (w / row_scale).round().clamp(-127, 127).to(torch.int8) + return q, (row_scale.view(-1) * scale).contiguous() + + +def quantize_w1(w1: torch.Tensor) -> tuple[torch.Tensor, ...]: + """U ~= s_hi * q_hi + s_lo * q_lo with the fixed ratio s_lo = s_hi / 256.""" + u = w1.float() + s_hi = (u.abs().amax(1, keepdim=True) / 127.0).clamp_min(1e-30) + q_hi = (u / s_hi).round().clamp(-127, 127) + s_lo = s_hi / 256.0 + q_lo = ((u - s_hi * q_hi) / s_lo).round().clamp(-127, 127) + return q_hi.to(torch.int8), q_lo.to(torch.int8), s_hi.view(-1), s_lo.view(-1) + + +def build_w2_tiles( + q_w2: torch.Tensor, rows_pad: int, tiles: int, mask: int +) -> tuple[torch.Tensor, torch.Tensor]: + """Per CTA, 64-row tiles as K-major wgmma B operands with permuted rows, + split into the resident and the streamed tiles (each [grid, n, 16384]).""" + qp = torch.zeros(rows_pad, MARKOV_RANK, dtype=torch.int8, device=q_w2.device) + qp[: q_w2.shape[0]] = q_w2 + num_tiles = rows_pad // TILE_ROWS + t = qp.view(num_tiles, TILE_ROWS, MARKOV_RANK)[:, _RHO] + # [tile][k32 chunk][8-column group][k16 core][column][16 B] + v = t.view(num_tiles, 8, 8, 8, 2, 16).permute(0, 3, 1, 4, 2, 5) + v = v.contiguous().view(num_tiles // tiles, tiles, TILE_ROWS * MARKOV_RANK) + streamed = torch.tensor( + [(mask >> i) & 1 for i in range(tiles)], dtype=torch.bool, device=qp.device + ) + return v[:, ~streamed].contiguous(), v[:, streamed].contiguous() + + +def build_w1_rows(q_hi, q_lo, s_hi, s_lo) -> torch.Tensor: + """W1 rows in A-fragment order: per k32 chunk j and lane tig, + q_hi[32j + 4tig ..] | q_lo[..] | q_hi[32j + 16 + 4tig ..] | q_lo[..], + then s_hi and s_lo as fp32 (528 B per row).""" + vocab = q_hi.shape[0] + hi = q_hi.view(vocab, 8, 2, 4, 4) + lo = q_lo.view(vocab, 8, 2, 4, 4) + frag = torch.stack([hi[:, :, 0], lo[:, :, 0], hi[:, :, 1], lo[:, :, 1]], dim=3) + rows = torch.zeros(vocab, W1_ROW_BYTES, dtype=torch.uint8, device=q_hi.device) + rows[:, :512] = frag.reshape(vocab, 512).view(torch.uint8) + rows[:, 512:520] = torch.stack([s_hi, s_lo], 1).contiguous().view(torch.uint8) + return rows + + +class DSparkMarkovWalker: + """int8 layouts of one Markov head and the kernel's persistent state.""" + + def __init__( + self, + markov_w1: torch.Tensor, + markov_w2: torch.Tensor, + num_steps: int, + logit_scale: float = 1.0, + ) -> None: + reason = unsupported_reason(markov_w1, markov_w2, num_steps) + if reason is not None: + raise ValueError(f"dspark_markov_quantization='int8': {reason}.") + device = markov_w2.device + self.vocab_size = markov_w2.shape[0] + self.num_steps = num_steps + num_sms = _num_sms(device) + tiles = pick_tiles(self.vocab_size, num_sms) + assert tiles is not None + self.tiles = tiles + self.mask = stream_mask(tiles) + rows_cta = tiles * TILE_ROWS + rows_pad = -(-self.vocab_size // rows_cta) * rows_cta + with torch.no_grad(): + q_w2, row_scale = quantize_w2(markov_w2, logit_scale) + self.w2_res, self.w2_str = build_w2_tiles(q_w2, rows_pad, tiles, self.mask) + del q_w2 + self.row_scale = torch.full( + (rows_pad,), 1e-30, dtype=torch.float32, device=device + ) + self.row_scale[: self.vocab_size] = row_scale + self.w1_rows = build_w1_rows(*quantize_w1(markov_w1)) + # Round counter and exchange slots; zeroed once, never reset. + self.state = torch.zeros(STATE_NUMEL, dtype=torch.int64, device=device) + assert self.state.data_ptr() % 128 == 0 + logger.info_once( + "DSpark Markov walk: int8 W2 on chip (%d tiles per CTA, %d CTAs, " + "%.1f MiB of layouts).", + tiles, + rows_pad // rows_cta, + self.memory_bytes() / 2**20, + ) + + def memory_bytes(self) -> int: + return sum( + t.numel() * t.element_size() + for t in (self.w2_res, self.w2_str, self.row_scale, self.w1_rows) + ) + + def supports(self, num_reqs: int, base_logits: torch.Tensor) -> bool: + return ( + 1 <= num_reqs <= MAX_REQS + and base_logits.dtype == torch.bfloat16 + and base_logits.shape[-1] == self.vocab_size + and base_logits.is_contiguous() + ) + + def walk( + self, + base_logits: torch.Tensor, + anchor_ids: torch.Tensor, + anchor_stride: int, + idx_mapping: torch.Tensor, + temperature: torch.Tensor | None, + seeds: torch.Tensor, + pos: torch.Tensor, + draft_tokens: torch.Tensor, + draft_logits: torch.Tensor | None, + ) -> None: + """Walk ``base_logits`` [num_reqs, num_steps, V] into ``draft_tokens``. + + No allocation and no host sync: safe under CUDA graph capture. + """ + ops.dspark_markov_walk( + self.w2_res, + self.w2_str, + self.row_scale, + self.w1_rows, + base_logits, + anchor_ids, + anchor_stride, + idx_mapping, + temperature, + seeds, + pos, + draft_tokens, + draft_logits, + self.state, + self.num_steps, + self.tiles, + self.mask, + ) + + def warmup(self) -> None: + """Load the kernel before CUDA graph capture. + + The first launch of a cooperative kernel loads its module; doing that + while another cooperative grid is spinning can deadlock. + """ + device = self.state.device + n = self.num_steps + base = torch.zeros(1, n, self.vocab_size, dtype=torch.bfloat16, device=device) + anchor = torch.zeros(1, dtype=torch.int32, device=device) + idx = torch.full((n,), -1, dtype=torch.int64, device=device) + seeds = torch.zeros(1, dtype=torch.int64, device=device) + pos = torch.zeros(n, dtype=torch.int64, device=device) + tokens = torch.zeros(1, n, dtype=torch.int64, device=device) + self.walk(base, anchor, 1, idx, None, seeds, pos, tokens, None) + torch.accelerator.synchronize(device) diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py index c29f9adce424..217bfffb944e 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/speculator.py @@ -32,6 +32,7 @@ from vllm.logger import init_logger from vllm.v1.worker.gpu.sample.gumbel import gumbel_sample from vllm.v1.worker.gpu.spec_decode.dflash.speculator import DFlashSpeculator +from vllm.v1.worker.gpu.spec_decode.dspark.markov_walk import DSparkMarkovWalker from vllm.v1.worker.gpu.spec_decode.dspark.utils import load_dspark_model logger = init_logger(__name__) @@ -80,6 +81,8 @@ def __init__(self, vllm_config: VllmConfig, device: torch.device): ) self.use_confidence_head: bool = False + # One-launch Markov walk (dspark_markov_quantization="int8"). + self._markov_walker: DSparkMarkovWalker | None = None def load_draft_model( self, @@ -111,8 +114,38 @@ def load_draft_model( # The acceptance estimator is not needed when a trained confidence head # is available. self.use_acceptance_estimator = False + if self.speculative_config.dspark_markov_quantization == "int8": + self._markov_walker = self._build_markov_walker(model) return model + def _build_markov_walker(self, model: torch.nn.Module) -> DSparkMarkovWalker: + def fail(reason: str) -> ValueError: + return ValueError(f"dspark_markov_quantization='int8': {reason}.") + + # Both read every step's logits, which the kernel never materializes. + if self.draft_watermarker is not None: + raise fail("draft watermarking is not supported") + if self.use_acceptance_estimator: + raise fail( + "adaptive verification needs a confidence head (the online " + "acceptance estimator is not supported)" + ) + head = getattr(getattr(model, "model", None), "markov_head", None) + logits_processor = getattr(model, "logits_processor", None) + if head is None or logits_processor is None: + raise fail(f"{type(model).__name__} has no vanilla Markov head") + if logits_processor.soft_cap is not None: + raise fail("a logit soft cap is not supported") + w2 = head.markov_w2.weight[: head.markov_w2.org_vocab_size] + walker = DSparkMarkovWalker( + head.markov_w1.weight, + w2, + self.num_speculative_steps, + logit_scale=logits_processor.scale, + ) + walker.warmup() + return walker + def _sample_logits( self, logits: torch.Tensor, @@ -176,6 +209,12 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: vocab_size = base_logits.shape[-1] base_logits = base_logits.view(num_reqs, n_spec, vocab_size) + if self._markov_walker is not None and self._markov_walker.supports( + num_reqs, base_logits + ): + self._walk_markov_fused(num_reqs, sample_hidden, base_logits) + return + idx_map = self.sample_idx_mapping[:num_sample].view(num_reqs, n_spec) sample_pos = self.sample_pos[:num_sample].view(num_reqs, n_spec) confidence_markov_embeds = [] @@ -206,6 +245,45 @@ def _sample_sequential(self, num_reqs: int, head_hidden: torch.Tensor) -> None: num_reqs, n_spec ) + def _walk_markov_fused( + self, + num_reqs: int, + sample_hidden: torch.Tensor, + base_logits: torch.Tensor, + ) -> None: + """All Markov steps in one launch (dspark_markov_quantization="int8").""" + assert self._markov_walker is not None + n_spec = self.num_speculative_steps + num_sample = num_reqs * n_spec + # The anchor of request i is input_ids[i * num_query_per_req]. Greedy + # drafting ignores the temperature; probabilistic drafting caches the + # logits of every sampled row for rejection sampling. + self._markov_walker.walk( + base_logits, + self.input_buffers.input_ids, + self.num_query_per_req, + self.sample_idx_mapping[:num_sample], + self.temperature if self.draft_logits is not None else None, + self.seeds, + self.sample_pos[:num_sample], + self.draft_tokens, + self.draft_logits, + ) + if self.use_confidence_head: + # Step i's Markov input is the token before it: the anchor, then + # the first n_spec - 1 draft tokens. + anchor = self.input_buffers.input_ids[self._anchor_idx[:num_reqs]] + prev = torch.cat( + [anchor.unsqueeze(1).long(), self.draft_tokens[:num_reqs, :-1]], + dim=1, + ) + confidence = self.model.compute_confidence( + sample_hidden, self.model.markov_embed(prev.flatten()) + ) + self.draft_token_confidence_probs[:num_reqs] = confidence.view( + num_reqs, n_spec + ) + def _sample_sequential_topk(self, num_reqs: int, head_hidden: torch.Tensor) -> None: """Apply the sequential Markov head only to top-k base-logit candidates. From 5af18a75ecf11924270797f051585c039aa2c586 Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Sun, 1 Nov 2026 22:00:57 -0500 Subject: [PATCH 2/3] [Perf][Spec Decode] DSpark Markov walk: read the winner after an acquire on the count A 16-B relaxed load is not one atomic snapshot of {key, count}, so a CTA could see the full count with a stale key. Poll the count with acquire, then load the key. Also reject FP8 Markov heads: their stored values are not the weights, and the int8 layouts do not apply their separate scales. Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com> --- csrc/dspark_markov_walk.cu | 32 ++++++++++--------- tests/kernels/test_dspark_markov_walk.py | 1 + .../gpu/spec_decode/dspark/markov_walk.py | 4 ++- 3 files changed, 21 insertions(+), 16 deletions(-) diff --git a/csrc/dspark_markov_walk.cu b/csrc/dspark_markov_walk.cu index b2d23481fc4b..90f87e57df20 100644 --- a/csrc/dspark_markov_walk.cu +++ b/csrc/dspark_markov_walk.cu @@ -54,10 +54,10 @@ // sees exactly the distribution a sampled token was drawn from. // * Each step ends with one global argmax per request: every CTA folds its // best (order-preserving fp32 value, inverted row) key into a {key, count} -// pair with atomicMax and a release add, then polls the pair. Ties go to -// the smallest row, as torch.argmax. The pairs of a call alternate between -// two sets by call parity, so nothing needs a memset between CUDA graph -// replays. +// pair with atomicMax and a release add, polls the count with acquire, then +// reads the key. Ties go to the smallest row, as torch.argmax. The pairs of +// a call alternate between two sets by call parity, so nothing needs a memset +// between CUDA graph replays. // * Sampling is exact up to floating-point error: inside each 8-row item the // row is drawn by inverse CDF, across items by Gumbel-max on the item's // log2 mass. The noise is Philox keyed by the request's seed and counted by @@ -120,11 +120,10 @@ __device__ __forceinline__ u64 ld_relaxed(const u64* p) { : "memory"); return v; } -// {key, count} in one 16-B load -__device__ __forceinline__ ulonglong2 ld_relaxed_v2(const u64* p) { - ulonglong2 v; - asm volatile("ld.relaxed.gpu.global.v2.u64 {%0, %1}, [%2];" - : "=l"(v.x), "=l"(v.y) +__device__ __forceinline__ u64 ld_acquire(const u64* p) { + u64 v; + asm volatile("ld.acquire.gpu.global.u64 %0, [%1];" + : "=l"(v) : "l"(p) : "memory"); return v; @@ -683,13 +682,16 @@ __global__ void __launch_bounds__(kThreads, 1) markov_walk_kernel( } // both requests of a lane polled in the same round trip const u64 n = gridDim.x; - ulonglong2 v0 = make_ulonglong2(0, lane < nb ? 0 : n); - ulonglong2 v1 = make_ulonglong2(0, lane + 32 < nb ? 0 : n); + u64 c0 = lane < nb ? 0 : n; + u64 c1 = lane + 32 < nb ? 0 : n; do { - if (v0.y < n) v0 = ld_relaxed_v2(pair(k, lane)); - if (v1.y < n) v1 = ld_relaxed_v2(pair(k, lane + 32)); - } while (v0.y < n || v1.y < n); - const u64 key[2] = {v0.x, v1.x}; + if (c0 < n) c0 = ld_acquire(pair(k, lane) + 1); + if (c1 < n) c1 = ld_acquire(pair(k, lane + 32) + 1); + } while (c0 < n || c1 < n); + // the acquires order every CTA's atomicMax before these loads + const u64 key[2] = { + lane < nb ? ld_relaxed(pair(k, lane)) : 0ull, + lane + 32 < nb ? ld_relaxed(pair(k, lane + 32)) : 0ull}; #pragma unroll for (int qd = 0; qd < 2; ++qd) { const int bq = lane + 32 * qd; diff --git a/tests/kernels/test_dspark_markov_walk.py b/tests/kernels/test_dspark_markov_walk.py index 8605077c516d..d3d319ae9f34 100644 --- a/tests/kernels/test_dspark_markov_walk.py +++ b/tests/kernels/test_dspark_markov_walk.py @@ -232,4 +232,5 @@ def test_rejects_unsupported_heads(): assert unsupported_reason(w1, w2[:, :128], 7) is not None assert unsupported_reason(w1[:2048], w2, 7) is not None assert unsupported_reason(w1, w2, 17) is not None + assert unsupported_reason(w1, w2.to(torch.float8_e4m3fn), 7) is not None assert STATE_NUMEL == 16 + 2 * 16 * 64 * 16 diff --git a/vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py b/vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py index 3354f665322d..3c040596f36f 100644 --- a/vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py +++ b/vllm/v1/worker/gpu/spec_decode/dspark/markov_walk.py @@ -20,6 +20,8 @@ logger = init_logger(__name__) MARKOV_RANK = 256 +# FP8 heads keep separate scales, so their stored values are not the weights. +WEIGHT_DTYPES = (torch.bfloat16, torch.float16, torch.float32) MAX_STEPS = 16 MAX_REQS = 64 TILE_ROWS = 64 @@ -54,7 +56,7 @@ def unsupported_reason( return f"it needs a rank-{MARKOV_RANK} Markov head" if w1.shape != w2.shape: return "it needs a full-vocabulary draft (no draft-to-target map)" - if not w2.dtype.is_floating_point or not w1.dtype.is_floating_point: + if w2.dtype not in WEIGHT_DTYPES or w1.dtype not in WEIGHT_DTYPES: return "it needs an unquantized Markov projection" if w2.shape[0] % 8 != 0: return "it needs a vocabulary size divisible by 8" From 276395c5e0278f497c90f017425c4dfad5159a62 Mon Sep 17 00:00:00 2001 From: Chen Cheng <85005591+ischencheng@users.noreply.github.com> Date: Sun, 1 Nov 2026 22:41:32 -0500 Subject: [PATCH 3/3] [Perf][Spec Decode] DSpark Markov walk: set the shared memory attribute per launch cudaFuncSetAttribute applies to the current device, so a process-wide flag left a second GPU without the shared memory opt-in. Signed-off-by: Chen Cheng <85005591+ischencheng@users.noreply.github.com> --- csrc/dspark_markov_walk.cu | 9 +++------ 1 file changed, 3 insertions(+), 6 deletions(-) diff --git a/csrc/dspark_markov_walk.cu b/csrc/dspark_markov_walk.cu index 90f87e57df20..708dd89bfa9d 100644 --- a/csrc/dspark_markov_walk.cu +++ b/csrc/dspark_markov_walk.cu @@ -757,12 +757,9 @@ void launch(const torch::stable::Tensor& w2_res, "base is too large for 32-bit offsets"); const int smem = (kRes + kRing) * kTileBytes + kRowsCta * 4; auto kernel = markov_walk_kernel; - static bool attr_set = false; - if (!attr_set) { - STD_CUDA_CHECK(cudaFuncSetAttribute( - kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); - attr_set = true; - } + // per device, so not cached across launches + STD_CUDA_CHECK(cudaFuncSetAttribute( + kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem)); cudaLaunchConfig_t cfg = {}; cfg.gridDim = dim3(grid); cfg.blockDim = dim3(kThreads);