From c3b2305a46f852d0d58dde685b711e6c1b7c172f Mon Sep 17 00:00:00 2001 From: DINMAY KUMAR BRAHMA Date: Fri, 10 Jul 2026 19:57:35 +0530 Subject: [PATCH] Add optimized DSpark CUDA primitives --- DSpark/README.md | 177 +++++++++++++ DSpark/__init__.py | 21 ++ DSpark/benchmark.py | 117 +++++++++ DSpark/csrc/bindings.cpp | 14 + DSpark/csrc/dspark_cuda.cu | 479 +++++++++++++++++++++++++++++++++++ DSpark/csrc/dspark_cuda.h | 16 ++ DSpark/dspark.py | 359 ++++++++++++++++++++++++++ DSpark/install.sh | 5 + DSpark/reference.py | 115 +++++++++ DSpark/setup.py | 44 ++++ DSpark/tests/test_dspark.py | 134 ++++++++++ MANIFEST.in | 10 +- README.md | 59 ++++- Usage/DSpark/dspark_usage.py | 48 ++++ Usage/DSpark/readme.md | 12 + pyproject.toml | 3 +- setup.py | 8 +- 17 files changed, 1615 insertions(+), 6 deletions(-) create mode 100644 DSpark/README.md create mode 100644 DSpark/__init__.py create mode 100644 DSpark/benchmark.py create mode 100644 DSpark/csrc/bindings.cpp create mode 100644 DSpark/csrc/dspark_cuda.cu create mode 100644 DSpark/csrc/dspark_cuda.h create mode 100644 DSpark/dspark.py create mode 100755 DSpark/install.sh create mode 100644 DSpark/reference.py create mode 100644 DSpark/setup.py create mode 100644 DSpark/tests/test_dspark.py create mode 100644 Usage/DSpark/dspark_usage.py create mode 100644 Usage/DSpark/readme.md diff --git a/DSpark/README.md b/DSpark/README.md new file mode 100644 index 0000000..6cd5d30 --- /dev/null +++ b/DSpark/README.md @@ -0,0 +1,177 @@ +# DSpark CUDA primitives + +Low-level PyTorch/CUDA inference primitives for **DSpark: Confidence-Scheduled +Speculative Decoding with Semi-Autoregressive Generation** (DeepSeek-AI, +2026). + +This module accelerates the two DSpark operations that sit directly on the +decode hot path: + +1. the default low-rank Markov head used to inject one-token local + autoregression into a parallel draft block; and +2. the hardware-aware prefix scheduler that decides how many draft tokens from + every active request should be sent to the target model. + +It is intentionally a kernel library, not a duplicate of DeepSeek's complete +DeepSpec training/evaluation stack. Use a trained DSpark drafter (or the +official checkpoints) to produce base logits and conditional confidence logits, +then pass those tensors to these primitives. + +## Kernel design + +### Markov logit correction + +For previous token `x`, the default head computes + +```text +corrected_logits = base_logits + W2 @ W1[x] +``` + +The CUDA implementation stores `W2.T` as `[rank, vocab]`, stages `W1[x]` in +shared memory, computes four coalesced vocabulary stripes per CTA, accumulates +with register FMAs, and fuses the final base-logit addition. It does not +allocate the intermediate vocabulary-sized Markov bias. The specialized path +is intended for latency-sensitive decode microbatches; training automatically +uses native PyTorch so autograd remains exact. + +### Confidence scheduler + +Given conditional acceptance estimates `c[r, k]`, calibrated prefix survival is + +```text +p[r, k] = product(sigmoid(logit[r, j] / temperature[j]), j=0..k) +``` + +The extension uses: + +- one warp per request and shuffle-based prefix products; +- a 64-bit key whose low word enforces prefix order on exact ties; +- CUB device radix sort and inclusive scan on the current PyTorch stream; +- a parallel first-throughput-drop search matching DSpark Algorithm 1; and +- device-side schedule scatter/finalization with no host synchronization. + +For `R` requests and `K` draft positions, candidate ranking is `O(RK)` radix +work and the auxiliary storage is `O(RK)`. `K` may be 1–32; the paper's default +is 7. + +## Install + +Requirements: + +- Linux with an NVIDIA GPU +- CUDA toolkit compatible with the installed PyTorch +- PyTorch 2.1 or newer +- a C++17 compiler + +```bash +cd DSpark +chmod +x install.sh +./install.sh +``` + +PyTorch chooses the local GPU architecture automatically. To build a portable +binary, set `TORCH_CUDA_ARCH_LIST`, for example: + +```bash +TORCH_CUDA_ARCH_LIST="6.1;7.5;8.0;8.6;8.9;9.0" ./install.sh +``` + +Compute capability 6.1 includes the GTX 1050 Ti. Newer architectures use the +same instruction-level path and benefit from larger caches and memory bandwidth. + +## Python API + +```python +import torch +from DSpark import DSparkMarkovHead, DSparkScheduler + +device = "cuda" +requests, proposal_length = 128, 7 +vocab_size, rank = 32_000, 256 + +# Semi-autoregressive Markov correction for one draft position. +head = DSparkMarkovHead(vocab_size, rank).to(device).eval() +base_logits = torch.randn(requests, vocab_size, device=device, dtype=torch.float16) +previous_ids = torch.randint(vocab_size, (requests,), device=device) +with torch.inference_mode(): + corrected_logits = head(base_logits, previous_ids) + +# Or correct and sample every position from parallel drafter logits [R, K, V]. +parallel_logits = torch.randn( + requests, proposal_length, vocab_size, device=device, dtype=torch.float16 +) +with torch.inference_mode(): + draft = head.sample_block(parallel_logits, previous_ids, temperature=0.0) + +# Confidence-scheduled verification. +confidence_logits = torch.randn( + requests, proposal_length, device=device, dtype=torch.float16 +) +max_physical_batch = requests * (proposal_length + 1) + +# step_curve[b] = profiled target-model steps/second at b verification tokens. +batch_sizes = torch.arange(max_physical_batch + 1, device=device) +step_curve = 900.0 / (1.0 + batch_sizes / 256.0) + +scheduler = DSparkScheduler( + proposal_length, + temperatures=[1.08, 1.04, 1.0, 0.98, 0.96, 0.95, 0.94], +).to(device) +result = scheduler(confidence_logits, step_curve) + +# Keep these tensors on-device when wiring them into a serving engine. +verification_lengths = result.lengths +expected_tokens = result.expected_tokens +expected_token_throughput = result.expected_throughput +``` + +`step_curve[batch_tokens]` must contain the profiled target-model steps per +second at that physical verification batch size. The scheduler includes one +anchor/bonus token per active request, so the table must cover indices through +`requests * (proposal_length + 1)`. + +If the extension is not built, both operations use equivalent PyTorch code. The +fallback is useful for CPU correctness checks and integration development. + +## Import official DeepSpec weights + +DeepSpec stores the second Markov matrix as `nn.Linear.weight` with shape +`[vocab, rank]`. This library stores the transpose for coalesced CUDA access: + +```python +head.load_deepspec_( + official_model.markov_head.markov_w1.weight, + official_model.markov_head.markov_w2.weight, +) +``` + +## Validate and benchmark + +From the repository root: + +```bash +python -m pytest DSpark/tests -q +python -m DSpark.benchmark --dtype float16 --requests 128 --vocab-size 32000 +``` + +The benchmark reports actual timings for the current GPU; this repository does +not claim a universal speedup because the Markov kernel's crossover against +cuBLAS depends on batch size, vocabulary, rank, dtype, and architecture. + +## Scope and guarantees + +- Inference forward path only for the custom kernels. +- Markov head: FP32, FP16, or BF16; rank up to 4096. +- Scheduler: FP32, FP16, or BF16 confidence logits; proposal length up to 32. +- All CUDA work is submitted to the active PyTorch stream and respects the + current CUDA device guard. +- CPU and autograd paths use native PyTorch operations. + +## References + +- [DSpark paper (arXiv:2607.05147)](https://arxiv.org/abs/2607.05147) +- [DeepSeek-AI/DeepSpec reference implementation](https://github.com/deepseek-ai/DeepSpec) + +The implementation in this directory is original MIT-licensed code written for +this repository and uses the public algorithm/tensor contracts from the sources +above. diff --git a/DSpark/__init__.py b/DSpark/__init__.py new file mode 100644 index 0000000..91ef924 --- /dev/null +++ b/DSpark/__init__.py @@ -0,0 +1,21 @@ +"""CUDA-accelerated primitives for DeepSeek DSpark.""" + +from .dspark import ( + DraftBlockResult, + DSparkMarkovHead, + DSparkScheduler, + ScheduleResult, + cuda_extension_available, + markov_logits, + schedule, +) + +__all__ = [ + "DraftBlockResult", + "DSparkMarkovHead", + "DSparkScheduler", + "ScheduleResult", + "cuda_extension_available", + "markov_logits", + "schedule", +] diff --git a/DSpark/benchmark.py b/DSpark/benchmark.py new file mode 100644 index 0000000..063e64c --- /dev/null +++ b/DSpark/benchmark.py @@ -0,0 +1,117 @@ +"""Microbenchmarks for the DSpark CUDA primitives.""" + +from __future__ import annotations + +import argparse +import statistics + +import torch + +from .dspark import cuda_extension_available, markov_logits, schedule +from .reference import markov_logits_reference, schedule_reference + + +DTYPES = { + "float16": torch.float16, + "bfloat16": torch.bfloat16, + "float32": torch.float32, +} + + +def _time_cuda(function, warmup: int, iterations: int) -> float: + for _ in range(warmup): + function() + torch.cuda.synchronize() + samples = [] + for _ in range(iterations): + start = torch.cuda.Event(enable_timing=True) + end = torch.cuda.Event(enable_timing=True) + start.record() + function() + end.record() + end.synchronize() + samples.append(start.elapsed_time(end) * 1_000.0) + return statistics.median(samples) + + +def main() -> None: + parser = argparse.ArgumentParser() + parser.add_argument("--requests", type=int, default=128) + parser.add_argument("--proposal-length", type=int, default=7) + parser.add_argument("--vocab-size", type=int, default=32_000) + parser.add_argument("--rank", type=int, default=256) + parser.add_argument("--dtype", choices=DTYPES, default="float16") + parser.add_argument("--warmup", type=int, default=20) + parser.add_argument("--iterations", type=int, default=100) + args = parser.parse_args() + + if not torch.cuda.is_available(): + raise SystemExit("CUDA is required for this benchmark") + if not cuda_extension_available(): + raise SystemExit("Build DSpark/install.sh before running the benchmark") + + device = torch.device("cuda") + dtype = DTYPES[args.dtype] + logits = torch.randn( + args.requests, + args.vocab_size, + device=device, + dtype=dtype, + ) + ids = torch.randint(args.vocab_size, (args.requests,), device=device) + embedding = torch.randn( + args.vocab_size, + args.rank, + device=device, + dtype=dtype, + ) + projection_t = torch.randn( + args.rank, + args.vocab_size, + device=device, + dtype=dtype, + ) + + confidence = torch.randn( + args.requests, + args.proposal_length, + device=device, + dtype=dtype, + ) + temperatures = torch.ones(args.proposal_length, device=device) + curve_size = args.requests * (args.proposal_length + 1) + 1 + curve_index = torch.arange(curve_size, device=device, dtype=torch.float32) + step_curve = 1_000.0 / (1.0 + curve_index / 256.0) + + with torch.inference_mode(): + custom_markov_us = _time_cuda( + lambda: markov_logits(logits, ids, embedding, projection_t), + args.warmup, + args.iterations, + ) + torch_markov_us = _time_cuda( + lambda: markov_logits_reference(logits, ids, embedding, projection_t), + args.warmup, + args.iterations, + ) + custom_schedule_us = _time_cuda( + lambda: schedule(confidence, step_curve, temperatures), + args.warmup, + args.iterations, + ) + torch_schedule_us = _time_cuda( + lambda: schedule_reference(confidence, step_curve, temperatures), + args.warmup, + args.iterations, + ) + + print(f"GPU: {torch.cuda.get_device_name()}") + print(f"dtype={args.dtype}, requests={args.requests}") + print(f"Markov CUDA: {custom_markov_us:9.2f} us") + print(f"Markov PyTorch: {torch_markov_us:9.2f} us") + print(f"Scheduler CUDA: {custom_schedule_us:9.2f} us") + print(f"Scheduler Torch: {torch_schedule_us:9.2f} us") + + +if __name__ == "__main__": + main() diff --git a/DSpark/csrc/bindings.cpp b/DSpark/csrc/bindings.cpp new file mode 100644 index 0000000..80136fa --- /dev/null +++ b/DSpark/csrc/bindings.cpp @@ -0,0 +1,14 @@ +#include + +#include "dspark_cuda.h" + +PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { + module.def( + "markov_logits", + &dspark_markov_logits_cuda, + "Fused DSpark low-rank Markov logit correction (CUDA)"); + module.def( + "schedule", + &dspark_schedule_cuda, + "DSpark confidence-scheduled verification (CUDA)"); +} diff --git a/DSpark/csrc/dspark_cuda.cu b/DSpark/csrc/dspark_cuda.cu new file mode 100644 index 0000000..acfa4cd --- /dev/null +++ b/DSpark/csrc/dspark_cuda.cu @@ -0,0 +1,479 @@ +#include + +#include +#include +#include +#include +#include +#include + +#include +#include +#include +#include + +#include "dspark_cuda.h" + +namespace { + +constexpr int kThreads = 256; +constexpr int kWarpSize = 32; +constexpr int kVocabItemsPerThread = 4; + +template +__device__ __forceinline__ float as_float(scalar_t value) { + return static_cast(value); +} + +// The default DSpark head is a low-rank Markov correction: +// logits[b, v] += W2[v, :] dot W1[previous_token[b], :]. +// W2 is supplied transposed as [rank, vocab], so every warp reads contiguous +// vocabulary lanes. The latent row is staged once per CTA and the vocabulary +// bias is never materialized. +template +__global__ __launch_bounds__(kThreads, 2) void markov_logits_kernel( + const scalar_t* __restrict__ base, + const int64_t* __restrict__ previous_ids, + const scalar_t* __restrict__ embedding, + const scalar_t* __restrict__ projection_t, + scalar_t* __restrict__ output, + int vocab_size, + int rank, + int vocab_tiles) { + extern __shared__ float latent[]; + + const int batch = blockIdx.x / vocab_tiles; + const int tile = blockIdx.x - batch * vocab_tiles; + const int64_t token = previous_ids[batch]; + + if (token < 0 || token >= vocab_size) { + const int tile_start = tile * blockDim.x * kVocabItemsPerThread; +#pragma unroll + for (int item = 0; item < kVocabItemsPerThread; ++item) { + const int vocab = tile_start + item * blockDim.x + threadIdx.x; + if (vocab < vocab_size) { + const int64_t offset = + static_cast(batch) * vocab_size + vocab; + output[offset] = base[offset]; + } + } + return; + } + + for (int r = threadIdx.x; r < rank; r += blockDim.x) { + latent[r] = as_float(embedding[token * rank + r]); + } + __syncthreads(); + + const int tile_start = tile * blockDim.x * kVocabItemsPerThread; + float correction[kVocabItemsPerThread] = {0.0f, 0.0f, 0.0f, 0.0f}; +#pragma unroll 4 + for (int r = 0; r < rank; ++r) { + const float latent_value = latent[r]; +#pragma unroll + for (int item = 0; item < kVocabItemsPerThread; ++item) { + const int vocab = tile_start + item * blockDim.x + threadIdx.x; + if (vocab < vocab_size) { + correction[item] = fmaf( + latent_value, + as_float(projection_t[ + static_cast(r) * vocab_size + vocab]), + correction[item]); + } + } + } + +#pragma unroll + for (int item = 0; item < kVocabItemsPerThread; ++item) { + const int vocab = tile_start + item * blockDim.x + threadIdx.x; + if (vocab < vocab_size) { + const int64_t offset = + static_cast(batch) * vocab_size + vocab; + output[offset] = + static_cast(as_float(base[offset]) + correction[item]); + } + } +} + +// One warp handles one request. DSpark blocks are deliberately short (7 in +// the paper); keeping the full calibrated prefix scan inside a warp removes +// shared-memory traffic and all inter-CTA synchronization. +template +__global__ __launch_bounds__(kThreads) void build_candidates_kernel( + const scalar_t* __restrict__ logits, + const float* __restrict__ temperatures, + float* __restrict__ survival, + uint64_t* __restrict__ keys, + int32_t* __restrict__ candidate_ids, + int request_count, + int proposal_length) { + const int lane = threadIdx.x & (kWarpSize - 1); + const int warp = threadIdx.x / kWarpSize; + const int warps_per_block = blockDim.x / kWarpSize; + const int request = blockIdx.x * warps_per_block + warp; + if (request >= request_count) { + return; + } + + float prefix = 1.0f; + if (lane < proposal_length) { + const int index = request * proposal_length + lane; + const float temperature = temperatures[lane]; + const float scaled_logit = as_float(logits[index]) / temperature; + prefix = 1.0f / (1.0f + __expf(-scaled_logit)); + } + +#pragma unroll + for (int offset = 1; offset < kWarpSize; offset <<= 1) { + const float previous = __shfl_up_sync(0xffffffffu, prefix, offset); + if (lane >= offset) { + prefix *= previous; + } + } + + if (lane < proposal_length) { + const int index = request * proposal_length + lane; + survival[index] = prefix; + + // Positive IEEE-754 floats preserve numeric order in their bit pattern. + // Position is the low-word tie break, so a prefix can never admit a + // later equal-probability token before its predecessor. + const uint64_t probability_bits = static_cast(__float_as_uint(prefix)); + const uint32_t position_tie_break = 0xffffffffu - static_cast(lane); + keys[index] = (probability_bits << 32) | position_tie_break; + candidate_ids[index] = index; + } +} + +__global__ __launch_bounds__(kThreads) void unpack_probabilities_kernel( + const uint64_t* __restrict__ sorted_keys, + float* __restrict__ sorted_probabilities, + int candidate_count) { + for (int index = blockIdx.x * blockDim.x + threadIdx.x; + index < candidate_count; + index += blockDim.x * gridDim.x) { + const uint32_t bits = static_cast(sorted_keys[index] >> 32); + sorted_probabilities[index] = __uint_as_float(bits); + } +} + +// A candidate is admitted only while it strictly improves expected token +// throughput. atomicMin finds the first drop, which is equivalent to the +// paper's causal early break but does not serialize the evaluation on one lane. +__global__ __launch_bounds__(kThreads) void find_first_drop_kernel( + const float* __restrict__ prefix_probability_sums, + const float* __restrict__ step_curve, + int32_t* __restrict__ first_drop, + int request_count, + int candidate_count) { + for (int index = blockIdx.x * blockDim.x + threadIdx.x; + index < candidate_count; + index += blockDim.x * gridDim.x) { + const float current_expected = + static_cast(request_count) + prefix_probability_sums[index]; + const int current_batch = request_count + index + 1; + const float current_throughput = current_expected * step_curve[current_batch]; + + float previous_expected = static_cast(request_count); + if (index > 0) { + previous_expected += prefix_probability_sums[index - 1]; + } + const int previous_batch = request_count + index; + const float previous_throughput = + previous_expected * step_curve[previous_batch]; + + if (!(current_throughput > previous_throughput)) { + atomicMin(first_drop, index); + } + } +} + +__global__ __launch_bounds__(kThreads) void scatter_schedule_kernel( + const int32_t* __restrict__ sorted_candidate_ids, + const int32_t* __restrict__ selected_count, + int32_t* __restrict__ lengths, + int candidate_count, + int proposal_length) { + const int count = *selected_count; + for (int index = blockIdx.x * blockDim.x + threadIdx.x; + index < candidate_count && index < count; + index += blockDim.x * gridDim.x) { + const int candidate = sorted_candidate_ids[index]; + const int request = candidate / proposal_length; + const int position = candidate - request * proposal_length; + atomicMax(lengths + request, position + 1); + } +} + +__global__ void finalize_schedule_kernel( + const float* __restrict__ prefix_probability_sums, + const float* __restrict__ step_curve, + const int32_t* __restrict__ selected_count, + float* __restrict__ expected_tokens, + float* __restrict__ expected_throughput, + int request_count) { + if (blockIdx.x != 0 || threadIdx.x != 0) { + return; + } + const int count = *selected_count; + float expected = static_cast(request_count); + if (count > 0) { + expected += prefix_probability_sums[count - 1]; + } + expected_tokens[0] = expected; + expected_throughput[0] = expected * step_curve[request_count + count]; +} + +inline int launch_blocks(int elements) { + return std::min((elements + kThreads - 1) / kThreads, 4096); +} + +void check_cuda_contiguous(const torch::Tensor& tensor, const char* name) { + TORCH_CHECK(tensor.is_cuda(), name, " must be a CUDA tensor"); + TORCH_CHECK(tensor.is_contiguous(), name, " must be contiguous"); +} + +} // namespace + +torch::Tensor dspark_markov_logits_cuda( + const torch::Tensor& base_logits, + const torch::Tensor& previous_token_ids, + const torch::Tensor& token_embedding, + const torch::Tensor& projection_t) { + check_cuda_contiguous(base_logits, "base_logits"); + check_cuda_contiguous(previous_token_ids, "previous_token_ids"); + check_cuda_contiguous(token_embedding, "token_embedding"); + check_cuda_contiguous(projection_t, "projection_t"); + + TORCH_CHECK(base_logits.dim() == 2, "base_logits must have shape [batch, vocab]"); + TORCH_CHECK(previous_token_ids.dim() == 1, "previous_token_ids must have shape [batch]"); + TORCH_CHECK(token_embedding.dim() == 2, "token_embedding must have shape [vocab, rank]"); + TORCH_CHECK(projection_t.dim() == 2, "projection_t must have shape [rank, vocab]"); + TORCH_CHECK(previous_token_ids.scalar_type() == torch::kInt64, + "previous_token_ids must be torch.int64"); + TORCH_CHECK(base_logits.scalar_type() == token_embedding.scalar_type() && + base_logits.scalar_type() == projection_t.scalar_type(), + "base_logits, token_embedding, and projection_t must have the same dtype"); + TORCH_CHECK(base_logits.device() == previous_token_ids.device() && + base_logits.device() == token_embedding.device() && + base_logits.device() == projection_t.device(), + "all inputs must be on the same CUDA device"); + TORCH_CHECK(base_logits.scalar_type() == torch::kFloat32 || + base_logits.scalar_type() == torch::kFloat16 || + base_logits.scalar_type() == torch::kBFloat16, + "Markov kernel supports float32, float16, and bfloat16"); + + const int64_t batch = base_logits.size(0); + const int64_t vocab = base_logits.size(1); + const int64_t rank = token_embedding.size(1); + TORCH_CHECK(batch > 0 && vocab > 0 && rank > 0, "input dimensions must be positive"); + TORCH_CHECK(batch <= std::numeric_limits::max() && + vocab <= std::numeric_limits::max() && + rank <= std::numeric_limits::max(), + "tensor dimensions exceed CUDA kernel limits"); + TORCH_CHECK(previous_token_ids.size(0) == batch, "batch dimension mismatch"); + TORCH_CHECK(token_embedding.size(0) == vocab, "embedding vocabulary mismatch"); + TORCH_CHECK(projection_t.size(0) == rank && projection_t.size(1) == vocab, + "projection_t must have shape [rank, vocab]"); + TORCH_CHECK(rank <= 4096, "rank > 4096 is not supported by the fused kernel"); + + c10::cuda::CUDAGuard device_guard(base_logits.device()); + const auto output = torch::empty_like(base_logits); + constexpr int vocab_values_per_cta = kThreads * kVocabItemsPerThread; + const int vocab_tiles = + (static_cast(vocab) + vocab_values_per_cta - 1) / + vocab_values_per_cta; + const int64_t grid_size = batch * vocab_tiles; + TORCH_CHECK(grid_size <= std::numeric_limits::max(), "launch grid is too large"); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + + AT_DISPATCH_FLOATING_TYPES_AND2( + torch::kFloat16, + torch::kBFloat16, + base_logits.scalar_type(), + "dspark_markov_logits_cuda", + [&] { + markov_logits_kernel<<< + static_cast(grid_size), + kThreads, + static_cast(rank) * sizeof(float), + stream>>>( + base_logits.data_ptr(), + previous_token_ids.data_ptr(), + token_embedding.data_ptr(), + projection_t.data_ptr(), + output.data_ptr(), + static_cast(vocab), + static_cast(rank), + vocab_tiles); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return output; +} + +std::vector dspark_schedule_cuda( + const torch::Tensor& confidence_logits, + const torch::Tensor& step_curve, + const torch::Tensor& temperatures) { + check_cuda_contiguous(confidence_logits, "confidence_logits"); + check_cuda_contiguous(step_curve, "step_curve"); + check_cuda_contiguous(temperatures, "temperatures"); + + TORCH_CHECK(confidence_logits.dim() == 2, + "confidence_logits must have shape [requests, proposal_length]"); + TORCH_CHECK(step_curve.dim() == 1, "step_curve must be one-dimensional"); + TORCH_CHECK(temperatures.dim() == 1, "temperatures must be one-dimensional"); + TORCH_CHECK(step_curve.scalar_type() == torch::kFloat32, + "step_curve must be torch.float32"); + TORCH_CHECK(temperatures.scalar_type() == torch::kFloat32, + "temperatures must be torch.float32"); + TORCH_CHECK(confidence_logits.scalar_type() == torch::kFloat32 || + confidence_logits.scalar_type() == torch::kFloat16 || + confidence_logits.scalar_type() == torch::kBFloat16, + "confidence_logits supports float32, float16, and bfloat16"); + TORCH_CHECK(confidence_logits.device() == step_curve.device() && + confidence_logits.device() == temperatures.device(), + "all inputs must be on the same CUDA device"); + + const int64_t request_count_64 = confidence_logits.size(0); + const int64_t proposal_length_64 = confidence_logits.size(1); + TORCH_CHECK(request_count_64 > 0, "at least one active request is required"); + TORCH_CHECK(proposal_length_64 > 0 && proposal_length_64 <= kWarpSize, + "proposal_length must be in [1, 32]"); + TORCH_CHECK(temperatures.size(0) == proposal_length_64, + "temperatures must contain one value per proposal position"); + + const int64_t candidate_count_64 = request_count_64 * proposal_length_64; + TORCH_CHECK(request_count_64 <= std::numeric_limits::max() && + candidate_count_64 <= std::numeric_limits::max(), + "candidate count exceeds CUDA/CUB limits"); + const int request_count = static_cast(request_count_64); + const int proposal_length = static_cast(proposal_length_64); + const int candidate_count = static_cast(candidate_count_64); + TORCH_CHECK(step_curve.size(0) > request_count + candidate_count, + "step_curve must be indexed through requests * (proposal_length + 1)"); + + c10::cuda::CUDAGuard device_guard(confidence_logits.device()); + const cudaStream_t stream = at::cuda::getCurrentCUDAStream(); + const auto float_options = confidence_logits.options().dtype(torch::kFloat32); + const auto int_options = confidence_logits.options().dtype(torch::kInt32); + const auto long_options = confidence_logits.options().dtype(torch::kInt64); + const auto byte_options = confidence_logits.options().dtype(torch::kUInt8); + + auto survival = torch::empty_like(confidence_logits, float_options); + auto keys_in = torch::empty({candidate_count}, long_options); + auto keys_out = torch::empty({candidate_count}, long_options); + auto ids_in = torch::empty({candidate_count}, int_options); + auto ids_out = torch::empty({candidate_count}, int_options); + + const int warps_per_block = kThreads / kWarpSize; + const int candidate_blocks = + (request_count + warps_per_block - 1) / warps_per_block; + AT_DISPATCH_FLOATING_TYPES_AND2( + torch::kFloat16, + torch::kBFloat16, + confidence_logits.scalar_type(), + "dspark_build_candidates_cuda", + [&] { + build_candidates_kernel<<>>( + confidence_logits.data_ptr(), + temperatures.data_ptr(), + survival.data_ptr(), + reinterpret_cast(keys_in.data_ptr()), + ids_in.data_ptr(), + request_count, + proposal_length); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + size_t sort_storage_bytes = 0; + C10_CUDA_CHECK(cub::DeviceRadixSort::SortPairsDescending( + nullptr, + sort_storage_bytes, + reinterpret_cast(keys_in.data_ptr()), + reinterpret_cast(keys_out.data_ptr()), + ids_in.data_ptr(), + ids_out.data_ptr(), + candidate_count, + 0, + 64, + stream)); + auto sort_storage = torch::empty( + {static_cast(sort_storage_bytes)}, byte_options); + C10_CUDA_CHECK(cub::DeviceRadixSort::SortPairsDescending( + sort_storage.data_ptr(), + sort_storage_bytes, + reinterpret_cast(keys_in.data_ptr()), + reinterpret_cast(keys_out.data_ptr()), + ids_in.data_ptr(), + ids_out.data_ptr(), + candidate_count, + 0, + 64, + stream)); + + auto sorted_probabilities = torch::empty({candidate_count}, float_options); + unpack_probabilities_kernel<<>>( + reinterpret_cast(keys_out.data_ptr()), + sorted_probabilities.data_ptr(), + candidate_count); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + auto prefix_probability_sums = torch::empty({candidate_count}, float_options); + size_t scan_storage_bytes = 0; + C10_CUDA_CHECK(cub::DeviceScan::InclusiveSum( + nullptr, + scan_storage_bytes, + sorted_probabilities.data_ptr(), + prefix_probability_sums.data_ptr(), + candidate_count, + stream)); + auto scan_storage = torch::empty( + {static_cast(scan_storage_bytes)}, byte_options); + C10_CUDA_CHECK(cub::DeviceScan::InclusiveSum( + scan_storage.data_ptr(), + scan_storage_bytes, + sorted_probabilities.data_ptr(), + prefix_probability_sums.data_ptr(), + candidate_count, + stream)); + + // candidate_count is both the no-drop sentinel and the selected count when + // every extension improves throughput. + auto selected_count = torch::full({1}, candidate_count, int_options); + find_first_drop_kernel<<>>( + prefix_probability_sums.data_ptr(), + step_curve.data_ptr(), + selected_count.data_ptr(), + request_count, + candidate_count); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + auto lengths = torch::zeros({request_count}, int_options); + scatter_schedule_kernel<<>>( + ids_out.data_ptr(), + selected_count.data_ptr(), + lengths.data_ptr(), + candidate_count, + proposal_length); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + auto expected_tokens = torch::empty({1}, float_options); + auto expected_throughput = torch::empty({1}, float_options); + finalize_schedule_kernel<<<1, 1, 0, stream>>>( + prefix_probability_sums.data_ptr(), + step_curve.data_ptr(), + selected_count.data_ptr(), + expected_tokens.data_ptr(), + expected_throughput.data_ptr(), + request_count); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + + return { + lengths, + survival, + selected_count, + expected_tokens, + expected_throughput, + }; +} diff --git a/DSpark/csrc/dspark_cuda.h b/DSpark/csrc/dspark_cuda.h new file mode 100644 index 0000000..13005c8 --- /dev/null +++ b/DSpark/csrc/dspark_cuda.h @@ -0,0 +1,16 @@ +#pragma once + +#include + +#include + +torch::Tensor dspark_markov_logits_cuda( + const torch::Tensor& base_logits, + const torch::Tensor& previous_token_ids, + const torch::Tensor& token_embedding, + const torch::Tensor& projection_t); + +std::vector dspark_schedule_cuda( + const torch::Tensor& confidence_logits, + const torch::Tensor& step_curve, + const torch::Tensor& temperatures); diff --git a/DSpark/dspark.py b/DSpark/dspark.py new file mode 100644 index 0000000..8a14547 --- /dev/null +++ b/DSpark/dspark.py @@ -0,0 +1,359 @@ +"""PyTorch interface for DeepSeek's DSpark inference primitives.""" + +from __future__ import annotations + +import math +from collections.abc import Sequence +from typing import NamedTuple + +import torch +from torch import nn + +from .reference import ( + ScheduleResult, + markov_logits_reference, + schedule_reference, +) + + +class DraftBlockResult(NamedTuple): + """Tokens and corrected logits from semi-autoregressive block sampling.""" + + token_ids: torch.Tensor + corrected_logits: torch.Tensor + + +try: + import _dspark_cuda +except ImportError: + _dspark_cuda = None + + +def cuda_extension_available() -> bool: + """Return whether the compiled CUDA extension was imported successfully.""" + + return _dspark_cuda is not None + + +def _check_markov_inputs( + base_logits: torch.Tensor, + previous_token_ids: torch.Tensor, + token_embedding: torch.Tensor, + projection_t: torch.Tensor, +) -> None: + if base_logits.ndim != 2: + raise ValueError("base_logits must have shape [batch, vocab]") + if previous_token_ids.ndim != 1: + raise ValueError("previous_token_ids must have shape [batch]") + if token_embedding.ndim != 2: + raise ValueError("token_embedding must have shape [vocab, rank]") + if projection_t.ndim != 2: + raise ValueError("projection_t must have shape [rank, vocab]") + + batch, vocab = base_logits.shape + embedding_vocab, rank = token_embedding.shape + if previous_token_ids.shape[0] != batch: + raise ValueError("previous_token_ids batch dimension does not match logits") + if embedding_vocab != vocab: + raise ValueError("token_embedding and base_logits vocabulary sizes differ") + if projection_t.shape != (rank, vocab): + raise ValueError("projection_t must have shape [rank, vocab]") + if previous_token_ids.dtype != torch.long: + raise TypeError("previous_token_ids must use torch.int64") + if not (base_logits.dtype == token_embedding.dtype == projection_t.dtype): + raise TypeError("base_logits, token_embedding, and projection_t need one dtype") + if not ( + base_logits.device + == previous_token_ids.device + == token_embedding.device + == projection_t.device + ): + raise ValueError("all Markov-head inputs must be on one device") + + +def markov_logits( + base_logits: torch.Tensor, + previous_token_ids: torch.Tensor, + token_embedding: torch.Tensor, + projection_t: torch.Tensor, +) -> torch.Tensor: + """Apply DSpark's default low-rank Markov head. + + On CUDA inference workloads this fuses embedding lookup, low-rank + projection, and base-logit addition. During autograd it intentionally uses + native PyTorch operations so parameter gradients remain exact. + """ + + _check_markov_inputs( + base_logits, + previous_token_ids, + token_embedding, + projection_t, + ) + needs_autograd = torch.is_grad_enabled() and any( + tensor.requires_grad + for tensor in (base_logits, token_embedding, projection_t) + ) + if base_logits.is_cuda and _dspark_cuda is not None and not needs_autograd: + return _dspark_cuda.markov_logits( + base_logits.contiguous(), + previous_token_ids.contiguous(), + token_embedding.contiguous(), + projection_t.contiguous(), + ) + return markov_logits_reference( + base_logits, + previous_token_ids, + token_embedding, + projection_t, + ) + + +def _temperatures_tensor( + temperatures: float | Sequence[float] | torch.Tensor | None, + *, + proposal_length: int, + device: torch.device, +) -> torch.Tensor: + if temperatures is None: + result = torch.ones(proposal_length, device=device, dtype=torch.float32) + elif isinstance(temperatures, torch.Tensor): + result = temperatures.to(device=device, dtype=torch.float32) + elif isinstance(temperatures, (float, int)): + result = torch.full( + (proposal_length,), + float(temperatures), + device=device, + dtype=torch.float32, + ) + else: + result = torch.as_tensor( + list(temperatures), + device=device, + dtype=torch.float32, + ) + if result.shape != (proposal_length,): + raise ValueError("temperatures must contain one value per proposal position") + positive = torch.all(result > 0) + if result.is_cuda: + torch._assert_async( + positive, + "all sequential calibration temperatures must be positive", + ) + elif not bool(positive.item()): + raise ValueError("all sequential calibration temperatures must be positive") + return result.contiguous() + + +@torch.no_grad() +def schedule( + confidence_logits: torch.Tensor, + step_curve: torch.Tensor | Sequence[float], + temperatures: float | Sequence[float] | torch.Tensor | None = None, +) -> ScheduleResult: + """Select per-request DSpark verification lengths. + + Args: + confidence_logits: Conditional acceptance logits with shape + ``[active_requests, proposal_length]``. + step_curve: Profiled target-model steps/second. Entry ``b`` is the + throughput for a physical verification batch of ``b`` tokens. It + must cover indices through ``requests * (proposal_length + 1)``. + temperatures: Optional STS calibration temperatures, one per proposal + position. A scalar applies the same temperature at every position. + + Returns: + A :class:`ScheduleResult`; ``lengths`` is the number of draft tokens to + verify for every request. Every returned tensor stays on the input + device, and the CUDA path does not synchronize with the host. + """ + + if confidence_logits.ndim != 2: + raise ValueError( + "confidence_logits must have shape [active_requests, proposal_length]" + ) + request_count, proposal_length = confidence_logits.shape + if request_count < 1: + raise ValueError("at least one active request is required") + if not 1 <= proposal_length <= 32: + raise ValueError("proposal_length must be in [1, 32]") + if confidence_logits.dtype not in (torch.float16, torch.bfloat16, torch.float32): + raise TypeError("confidence_logits must be float16, bfloat16, or float32") + + curve = torch.as_tensor( + step_curve, + dtype=torch.float32, + device=confidence_logits.device, + ).contiguous() + if curve.ndim != 1: + raise ValueError("step_curve must be one-dimensional") + required_curve_size = request_count * (proposal_length + 1) + 1 + if curve.numel() < required_curve_size: + raise ValueError( + f"step_curve needs at least {required_curve_size} entries; " + f"received {curve.numel()}" + ) + calibrated_temperatures = _temperatures_tensor( + temperatures, + proposal_length=proposal_length, + device=confidence_logits.device, + ) + + if confidence_logits.is_cuda and _dspark_cuda is not None: + tensors = _dspark_cuda.schedule( + confidence_logits.contiguous(), + curve, + calibrated_temperatures, + ) + return ScheduleResult(*tensors) + return schedule_reference( + confidence_logits, + curve, + calibrated_temperatures, + ) + + +class DSparkScheduler(nn.Module): + """Reusable module holding DSpark's sequential calibration temperatures.""" + + def __init__( + self, + proposal_length: int = 7, + temperatures: float | Sequence[float] | torch.Tensor | None = None, + ) -> None: + super().__init__() + if not 1 <= proposal_length <= 32: + raise ValueError("proposal_length must be in [1, 32]") + values = _temperatures_tensor( + temperatures, + proposal_length=proposal_length, + device=torch.device("cpu"), + ) + self.proposal_length = proposal_length + self.register_buffer("temperatures", values) + + def forward( + self, + confidence_logits: torch.Tensor, + step_curve: torch.Tensor | Sequence[float], + ) -> ScheduleResult: + if confidence_logits.shape[-1] != self.proposal_length: + raise ValueError( + f"expected proposal length {self.proposal_length}, " + f"received {confidence_logits.shape[-1]}" + ) + return schedule(confidence_logits, step_curve, self.temperatures) + + +class DSparkMarkovHead(nn.Module): + """Default semi-autoregressive DSpark Markov head. + + The projection is stored as ``[rank, vocab]``β€”the transpose of + ``nn.Linear.weight``β€”to make vocabulary lanes contiguous in the CUDA + kernel. Use :meth:`load_deepspec_` to import an official DeepSpec head. + """ + + def __init__(self, vocab_size: int, rank: int = 256) -> None: + super().__init__() + if vocab_size < 1 or rank < 1: + raise ValueError("vocab_size and rank must be positive") + self.vocab_size = int(vocab_size) + self.rank = int(rank) + self.token_embedding = nn.Parameter(torch.empty(vocab_size, rank)) + self.projection_t = nn.Parameter(torch.empty(rank, vocab_size)) + self.reset_parameters() + + def reset_parameters(self) -> None: + nn.init.normal_(self.token_embedding, std=1.0 / math.sqrt(self.rank)) + nn.init.xavier_uniform_(self.projection_t) + + @torch.no_grad() + def load_deepspec_( + self, + markov_w1_weight: torch.Tensor, + markov_w2_weight: torch.Tensor, + ) -> "DSparkMarkovHead": + """Load ``markov_w1.weight`` and ``markov_w2.weight`` from DeepSpec.""" + + if markov_w1_weight.shape != self.token_embedding.shape: + raise ValueError("markov_w1 weight shape does not match this head") + if markov_w2_weight.shape != (self.vocab_size, self.rank): + raise ValueError("markov_w2 weight must have shape [vocab, rank]") + self.token_embedding.copy_(markov_w1_weight) + self.projection_t.copy_(markov_w2_weight.t()) + return self + + def forward( + self, + base_logits: torch.Tensor, + previous_token_ids: torch.Tensor, + ) -> torch.Tensor: + return markov_logits( + base_logits, + previous_token_ids, + self.token_embedding, + self.projection_t, + ) + + @torch.no_grad() + def sample_block( + self, + base_logits: torch.Tensor, + first_previous_token_ids: torch.Tensor, + *, + temperature: float = 0.0, + generator: torch.Generator | None = None, + ) -> DraftBlockResult: + """Sample a DSpark block left-to-right through the Markov head. + + ``base_logits`` is the parallel drafter output with shape + ``[batch, proposal_length, vocab]``. A non-positive temperature uses + greedy decoding; a positive value samples from the corrected + distribution just like the reference DeepSpec loop. + """ + + if base_logits.ndim != 3 or base_logits.shape[-1] != self.vocab_size: + raise ValueError( + "base_logits must have shape [batch, proposal_length, vocab_size]" + ) + if first_previous_token_ids.shape != (base_logits.shape[0],): + raise ValueError("first_previous_token_ids must have shape [batch]") + if first_previous_token_ids.dtype != torch.long: + raise TypeError("first_previous_token_ids must use torch.int64") + + proposal_length = base_logits.shape[1] + if proposal_length == 0: + return DraftBlockResult( + torch.empty( + base_logits.shape[0], + 0, + dtype=torch.long, + device=base_logits.device, + ), + base_logits, + ) + + previous = first_previous_token_ids + sampled = [] + corrected = [] + for position in range(proposal_length): + position_logits = self(base_logits[:, position, :], previous) + corrected.append(position_logits) + if temperature <= 0.0: + next_tokens = position_logits.argmax(dim=-1) + else: + probabilities = torch.softmax( + position_logits.float() / float(temperature), + dim=-1, + ) + next_tokens = torch.multinomial( + probabilities, + num_samples=1, + generator=generator, + ).squeeze(-1) + sampled.append(next_tokens) + previous = next_tokens + return DraftBlockResult( + torch.stack(sampled, dim=1), + torch.stack(corrected, dim=1), + ) diff --git a/DSpark/install.sh b/DSpark/install.sh new file mode 100755 index 0000000..70ebe9c --- /dev/null +++ b/DSpark/install.sh @@ -0,0 +1,5 @@ +#!/usr/bin/env bash +set -euo pipefail + +cd "$(dirname "${BASH_SOURCE[0]}")" +python -m pip install --no-build-isolation -e . diff --git a/DSpark/reference.py b/DSpark/reference.py new file mode 100644 index 0000000..ae1eb75 --- /dev/null +++ b/DSpark/reference.py @@ -0,0 +1,115 @@ +"""Exact PyTorch reference operations for the DSpark CUDA extension.""" + +from __future__ import annotations + +from typing import NamedTuple + +import torch + + +class ScheduleResult(NamedTuple): + """Result of one hardware-aware DSpark scheduling decision.""" + + lengths: torch.Tensor + survival: torch.Tensor + selected_count: torch.Tensor + expected_tokens: torch.Tensor + expected_throughput: torch.Tensor + + +def markov_logits_reference( + base_logits: torch.Tensor, + previous_token_ids: torch.Tensor, + token_embedding: torch.Tensor, + projection_t: torch.Tensor, +) -> torch.Tensor: + """Apply the default DSpark low-rank Markov correction with PyTorch ops.""" + + latent = token_embedding[previous_token_ids.long()] + return base_logits + latent @ projection_t + + +def schedule_reference( + confidence_logits: torch.Tensor, + step_curve: torch.Tensor, + temperatures: torch.Tensor, +) -> ScheduleResult: + """Vectorized reference for Algorithm 1 in the DSpark paper. + + ``step_curve[b]`` is the profiled target-model steps/second for a physical + verification batch containing ``b`` tokens. One anchor/bonus token per + request is included before draft candidates are admitted. + """ + + request_count, proposal_length = confidence_logits.shape + conditional = torch.sigmoid(confidence_logits.float() / temperatures) + survival = conditional.cumprod(dim=-1) + + flat_survival = survival.flatten() + # Stable order plus row-major flattening keeps an earlier position ahead of + # a later equal-probability position from the same request. + order = torch.argsort(flat_survival, descending=True, stable=True) + sorted_survival = flat_survival[order] + prefix_sums = sorted_survival.cumsum(dim=0) + + candidate_count = flat_survival.numel() + candidate_index = torch.arange( + 1, + candidate_count + 1, + device=confidence_logits.device, + dtype=torch.long, + ) + expected = float(request_count) + prefix_sums + throughputs = expected * step_curve[request_count + candidate_index] + baseline = step_curve.new_tensor(float(request_count)) * step_curve[request_count] + previous = torch.cat([baseline.reshape(1), throughputs[:-1]]) + drops = ~(throughputs > previous) + + first_drop = torch.argmax(drops.to(torch.int64)) + selected_count_long = torch.where( + drops.any(), + first_drop, + first_drop.new_tensor(candidate_count), + ) + selected_count = selected_count_long.to(torch.int32).reshape(1) + + selected_mask = torch.arange( + candidate_count, + device=confidence_logits.device, + ) < selected_count_long + selected_ids = order[selected_mask] + selected_requests = torch.div( + selected_ids, + proposal_length, + rounding_mode="floor", + ) + selected_positions = selected_ids.remainder(proposal_length).to(torch.int32) + 1 + lengths = torch.zeros( + request_count, + dtype=torch.int32, + device=confidence_logits.device, + ) + lengths.scatter_reduce_( + 0, + selected_requests, + selected_positions, + reduce="amax", + include_self=True, + ) + + selected_prefix = prefix_sums[(selected_count_long - 1).clamp_min(0)] + final_expected = torch.where( + selected_count_long > 0, + selected_prefix + float(request_count), + prefix_sums.new_tensor(float(request_count)), + ).reshape(1) + final_throughput = ( + final_expected * step_curve[request_count + selected_count_long] + ).reshape(1) + return ScheduleResult( + lengths, + survival, + selected_count, + final_expected, + final_throughput, + ) diff --git a/DSpark/setup.py b/DSpark/setup.py new file mode 100644 index 0000000..91d9dbc --- /dev/null +++ b/DSpark/setup.py @@ -0,0 +1,44 @@ +"""Build the DSpark PyTorch CUDA extension.""" + +from pathlib import Path + +from setuptools import setup +from torch.utils.cpp_extension import BuildExtension, CUDAExtension + + +ROOT = Path(__file__).parent + +setup( + name="cuda-ml-dspark", + version="0.1.0", + description="Low-level CUDA inference primitives for DeepSeek DSpark", + long_description=(ROOT / "README.md").read_text(encoding="utf-8"), + long_description_content_type="text/markdown", + packages=["DSpark"], + package_dir={"DSpark": "."}, + ext_modules=[ + CUDAExtension( + name="_dspark_cuda", + sources=[ + str(ROOT / "csrc" / "bindings.cpp"), + str(ROOT / "csrc" / "dspark_cuda.cu"), + ], + include_dirs=[str(ROOT / "csrc")], + extra_compile_args={ + "cxx": ["-O3", "-std=c++17"], + "nvcc": [ + "-O3", + "-std=c++17", + "--use_fast_math", + "--expt-relaxed-constexpr", + "--expt-extended-lambda", + "-lineinfo", + ], + }, + ) + ], + cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)}, + python_requires=">=3.9", + install_requires=["torch>=2.1"], + zip_safe=False, +) diff --git a/DSpark/tests/test_dspark.py b/DSpark/tests/test_dspark.py new file mode 100644 index 0000000..58d8036 --- /dev/null +++ b/DSpark/tests/test_dspark.py @@ -0,0 +1,134 @@ +from __future__ import annotations + +import pytest +import torch + +from DSpark import ( + DSparkMarkovHead, + cuda_extension_available, + markov_logits, + schedule, +) +from DSpark.reference import markov_logits_reference, schedule_reference + + +def _curve(values: dict[int, float], size: int = 9) -> torch.Tensor: + curve = torch.zeros(size, dtype=torch.float32) + for index, value in values.items(): + curve[index] = value + return curve + + +def test_markov_reference_matches_formula_and_backpropagates() -> None: + torch.manual_seed(4) + base = torch.randn(3, 11, requires_grad=True) + ids = torch.tensor([2, 7, 4]) + embedding = torch.randn(11, 5, requires_grad=True) + projection_t = torch.randn(5, 11, requires_grad=True) + + actual = markov_logits(base, ids, embedding, projection_t) + expected = base + embedding[ids] @ projection_t + torch.testing.assert_close(actual, expected) + actual.square().mean().backward() + assert base.grad is not None + assert embedding.grad is not None + assert projection_t.grad is not None + + +def test_scheduler_stops_at_first_throughput_drop() -> None: + probability = torch.tensor(0.9) + logits = torch.logit(probability).expand(2, 3).clone() + curve = _curve({2: 100.0, 3: 95.0, 4: 90.0, 5: 70.0, 6: 60.0, 7: 50.0, 8: 40.0}) + + result = schedule(logits, curve) + + assert result.selected_count.item() == 2 + assert result.lengths.tolist() == [1, 1] + torch.testing.assert_close( + result.survival, + torch.tensor([[0.9, 0.81, 0.729], [0.9, 0.81, 0.729]]), + ) + torch.testing.assert_close(result.expected_tokens, torch.tensor([3.8])) + torch.testing.assert_close(result.expected_throughput, torch.tensor([342.0])) + + +def test_scheduler_can_select_none_or_everything() -> None: + logits = torch.zeros(2, 3) + none_curve = _curve({2: 100.0, 3: 10.0, 4: 9.0, 5: 8.0, 6: 7.0, 7: 6.0, 8: 5.0}) + none = schedule(logits, none_curve) + assert none.selected_count.item() == 0 + assert none.lengths.tolist() == [0, 0] + + all_curve = torch.ones(9) + all_candidates = schedule(logits, all_curve) + assert all_candidates.selected_count.item() == 6 + assert all_candidates.lengths.tolist() == [3, 3] + + +def test_sequential_temperature_scaling_is_applied_before_cumprod() -> None: + logits = torch.tensor([[2.0, 2.0]]) + curve = torch.ones(4) + result = schedule(logits, curve, temperatures=[2.0, 1.0]) + expected = torch.tensor( + [[torch.sigmoid(torch.tensor(1.0)), + torch.sigmoid(torch.tensor(1.0)) * torch.sigmoid(torch.tensor(2.0))]] + ) + torch.testing.assert_close(result.survival, expected) + + +def test_deepspec_weight_loader_transposes_projection() -> None: + head = DSparkMarkovHead(vocab_size=13, rank=4) + w1 = torch.randn(13, 4) + w2 = torch.randn(13, 4) + head.load_deepspec_(w1, w2) + torch.testing.assert_close(head.token_embedding, w1) + torch.testing.assert_close(head.projection_t, w2.t()) + + +def test_markov_head_samples_block_sequentially() -> None: + head = DSparkMarkovHead(vocab_size=5, rank=2) + with torch.no_grad(): + head.token_embedding.zero_() + head.projection_t.zero_() + base = torch.tensor( + [[[0.0, 3.0, 1.0, 0.0, -1.0], [0.0, 1.0, 4.0, 0.0, -1.0]]] + ) + result = head.sample_block(base, torch.tensor([0])) + assert result.token_ids.tolist() == [[1, 2]] + torch.testing.assert_close(result.corrected_logits, base) + + +@pytest.mark.skipif( + not torch.cuda.is_available() or not cuda_extension_available(), + reason="compiled DSpark CUDA extension is unavailable", +) +@pytest.mark.parametrize("dtype", [torch.float16, torch.float32]) +def test_cuda_matches_reference(dtype: torch.dtype) -> None: + torch.manual_seed(5) + device = torch.device("cuda") + base = torch.randn(4, 257, device=device, dtype=dtype) + ids = torch.randint(257, (4,), device=device) + embedding = torch.randn(257, 16, device=device, dtype=dtype) + projection_t = torch.randn(16, 257, device=device, dtype=dtype) + with torch.inference_mode(): + actual = markov_logits(base, ids, embedding, projection_t) + expected = markov_logits_reference(base, ids, embedding, projection_t) + tolerance = 2e-2 if dtype == torch.float16 else 2e-5 + torch.testing.assert_close(actual, expected, atol=tolerance, rtol=tolerance) + + confidence = torch.randn(17, 7, device=device, dtype=dtype) + temperatures = torch.ones(7, device=device) + curve = torch.ones(17 * 8 + 1, device=device) + actual_schedule = schedule(confidence, curve, temperatures) + expected_schedule = schedule_reference(confidence, curve, temperatures) + torch.testing.assert_close( + actual_schedule.survival, + expected_schedule.survival, + atol=2e-3, + rtol=2e-3, + ) + torch.testing.assert_close(actual_schedule.lengths, expected_schedule.lengths) + torch.testing.assert_close( + actual_schedule.selected_count, + expected_schedule.selected_count, + ) diff --git a/MANIFEST.in b/MANIFEST.in index 4765a2f..3973a9f 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -19,6 +19,14 @@ include HBM_SVM/*.cpp # Include build and test scripts include HBM_SVM/build_and_test.sh +# Include the DSpark PyTorch/CUDA extension and documentation +recursive-include DSpark *.py +recursive-include DSpark *.cpp +recursive-include DSpark *.cu +recursive-include DSpark *.h +recursive-include DSpark *.md +include DSpark/install.sh + # Include usage examples and documentation recursive-include Usage *.py recursive-include Usage *.md @@ -49,4 +57,4 @@ global-exclude .vscode global-exclude .idea global-exclude *.swp global-exclude *.swo -global-exclude *~ \ No newline at end of file +global-exclude *~ diff --git a/README.md b/README.md index 58830e6..449e075 100644 --- a/README.md +++ b/README.md @@ -1,6 +1,6 @@ # CUDA ML Library [![zread](https://img.shields.io/badge/Ask_Zread-_.svg?style=plastic&color=00b0aa&labelColor=000000&logo=data%3Aimage%2Fsvg%2Bxml%3Bbase64%2CPHN2ZyB3aWR0aD0iMTYiIGhlaWdodD0iMTYiIHZpZXdCb3g9IjAgMCAxNiAxNiIgZmlsbD0ibm9uZSIgeG1sbnM9Imh0dHA6Ly93d3cudzMub3JnLzIwMDAvc3ZnIj4KPHBhdGggZD0iTTQuOTYxNTYgMS42MDAxSDIuMjQxNTZDMS44ODgxIDEuNjAwMSAxLjYwMTU2IDEuODg2NjQgMS42MDE1NiAyLjI0MDFWNC45NjAxQzEuNjAxNTYgNS4zMTM1NiAxLjg4ODEgNS42MDAxIDIuMjQxNTYgNS42MDAxSDQuOTYxNTZDNS4zMTUwMiA1LjYwMDEgNS42MDE1NiA1LjMxMzU2IDUuNjAxNTYgNC45NjAxVjIuMjQwMUM1LjYwMTU2IDEuODg2NjQgNS4zMTUwMiAxLjYwMDEgNC45NjE1NiAxLjYwMDFaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik00Ljk2MTU2IDEwLjM5OTlIMi4yNDE1NkMxLjg4ODEgMTAuMzk5OSAxLjYwMTU2IDEwLjY4NjQgMS42MDE1NiAxMS4wMzk5VjEzLjc1OTlDMS42MDE1NiAxNC4xMTM0IDEuODg4MSAxNC4zOTk5IDIuMjQxNTYgMTQuMzk5OUg0Ljk2MTU2QzUuMzE1MDIgMTQuMzk5OSA1LjYwMTU2IDE0LjExMzQgNS42MDE1NiAxMy43NTk5VjExLjAzOTlDNS42MDE1NiAxMC42ODY0IDUuMzE1MDIgMTAuMzk5OSA0Ljk2MTU2IDEwLjM5OTlaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik0xMy43NTg0IDEuNjAwMUgxMS4wMzg0QzEwLjY4NSAxLjYwMDEgMTAuMzk4NCAxLjg4NjY0IDEwLjM5ODQgMi4yNDAxVjQuOTYwMUMxMC4zOTg0IDUuMzEzNTYgMTAuNjg1IDUuNjAwMSAxMS4wMzg0IDUuNjAwMUgxMy43NTg0QzE0LjExMTkgNS42MDAxIDE0LjM5ODQgNS4zMTM1NiAxNC4zOTg0IDQuOTYwMVYyLjI0MDFDMTQuMzk4NCAxLjg4NjY0IDE0LjExMTkgMS42MDAxIDEzLjc1ODQgMS42MDAxWiIgZmlsbD0iI2ZmZiIvPgo8cGF0aCBkPSJNNCAxMkwxMiA0TDQgMTJaIiBmaWxsPSIjZmZmIi8%2BCjxwYXRoIGQ9Ik00IDEyTDEyIDQiIHN0cm9rZT0iI2ZmZiIgc3Ryb2tlLXdpZHRoPSIxLjUiIHN0cm9rZS1saW5lY2FwPSJyb3VuZCIvPgo8L3N2Zz4K&logoColor=ffffff)](https://zread.ai/dino65-dev/Cuda_ML_Library.git) -A high-performance CUDA-accelerated Machine Learning library with automatic CPU fallback support, featuring optimized Support Vector Machine implementations for both classification and regression tasks. +A high-performance CUDA-accelerated machine-learning kernel library with Python integration and CPU/PyTorch fallbacks. ## πŸš€ Features @@ -11,6 +11,7 @@ A high-performance CUDA-accelerated Machine Learning library with automatic CPU - **Multiple Kernel Functions**: Linear, RBF, Polynomial, and Sigmoid kernels - **Advanced Algorithms**: SMO (Sequential Minimal Optimization) algorithm implementation - **FlashAttention**: Memory-efficient O(N) attention mechanism for transformer models with full training support +- **DeepSeek DSpark**: Fused low-rank Markov logits and confidence-scheduled speculative verification - **Memory Optimization**: Efficient GPU memory management with pooling - **Easy Integration**: Scikit-learn compatible API and PyTorch integration @@ -24,7 +25,7 @@ A high-performance CUDA-accelerated Machine Learning library with automatic CPU ### Software Requirements - **CUDA Toolkit** (Optional): Version 12.0+ for GPU acceleration - **Python**: 3.8+ -- **Dependencies**: numpy β‰₯1.19.0, scikit-learn β‰₯1.0.0 +- **Dependencies**: numpy β‰₯1.19.0, scikit-learn β‰₯1.0.0; PyTorch β‰₯2.1 for DSpark ### Supported Environments - **GPU-Accelerated**: Systems with CUDA-capable NVIDIA GPUs @@ -60,6 +61,13 @@ cd .. pip install -e . ``` +Build the optional DSpark PyTorch extension separately: + +```bash +cd DSpark +./install.sh +``` + The build process will: - Auto-detect CUDA availability and GPU architecture - Compile CUDA kernels when GPU is available @@ -137,6 +145,38 @@ print(f"Output shape: {output.shape}") # [2, 8, 512, 64] print(f"Memory efficient: O(N) instead of O(NΒ²)") ``` +### DeepSeek DSpark Example + +```python +import torch +from DSpark import DSparkMarkovHead, DSparkScheduler + +requests, proposal_length = 128, 7 +vocab_size, rank = 32_000, 256 + +head = DSparkMarkovHead(vocab_size, rank).cuda().eval() +scheduler = DSparkScheduler(proposal_length).cuda() + +base_logits = torch.randn(requests, vocab_size, device="cuda", dtype=torch.float16) +previous_ids = torch.randint(vocab_size, (requests,), device="cuda") +confidence_logits = torch.randn( + requests, proposal_length, device="cuda", dtype=torch.float16 +) + +max_batch = requests * (proposal_length + 1) +batch_tokens = torch.arange(max_batch + 1, device="cuda") +step_curve = 1_000.0 / (1.0 + batch_tokens / 256.0) + +with torch.inference_mode(): + corrected_logits = head(base_logits, previous_ids) + decision = scheduler(confidence_logits, step_curve) + +verification_lengths = decision.lengths +``` + +See the [DSpark CUDA documentation](./DSpark/README.md) for the kernel design, +DeepSpec weight import, step-curve contract, tests, and GPU benchmark. + ## πŸ“š API Reference ### CudaSVC (Classification) @@ -196,6 +236,19 @@ flash_attention( - Numerical accuracy < 1e-6 vs standard attention - Works with all PyTorch optimizers (Adam, SGD, etc.) +### DSpark + +```python +DSparkMarkovHead(vocab_size, rank=256) +DSparkScheduler(proposal_length=7, temperatures=None) +``` + +- FP32, FP16, and BF16 inference kernels +- Warp-level calibrated prefix products and CUB candidate ranking +- Paper-compatible first-throughput-drop admission rule +- No host synchronization on the CUDA scheduling path +- Native PyTorch fallback for CPU execution and Markov-head autograd + ## πŸ”§ Advanced Usage ### Hardware Detection @@ -234,6 +287,7 @@ svc_sigmoid = CudaSVC(kernel='sigmoid', gamma='scale', coef0=0.0) - **SVM**: Fully functional and ready for production use - **RF**: Fully functional and ready for production use - **FlashAttention**: Fully functional for training and inference (head_dim=64, FP32 only) +- **DSpark**: CUDA inference primitives plus exact PyTorch fallback; integrate with a trained DSpark drafter **Note**: For production transformer workloads with advanced features (FP16, variable head dimensions, attention masks), consider using the official [FlashAttention](https://github.com/Dao-AILab/flash-attention) implementation. This implementation is ideal for learning, prototyping, and small-scale training. @@ -367,6 +421,7 @@ This project is licensed under the MIT License - see the [LICENSE](LICENSE) file - [SVM Usage Examples](./Usage/SVM/) - [Random Forest Usage Examples](./Usage/Random_forest/) - [FlashAttention Documentation](./flash_attention/USAGE.md) + - [DSpark CUDA Documentation](./DSpark/README.md) ## πŸ“Š Version diff --git a/Usage/DSpark/dspark_usage.py b/Usage/DSpark/dspark_usage.py new file mode 100644 index 0000000..0660a0e --- /dev/null +++ b/Usage/DSpark/dspark_usage.py @@ -0,0 +1,48 @@ +"""Minimal DSpark CUDA integration example.""" + +import torch + +from DSpark import DSparkMarkovHead, DSparkScheduler + + +def main() -> None: + if not torch.cuda.is_available(): + raise SystemExit("This example requires an NVIDIA GPU") + + device = torch.device("cuda") + requests, proposal_length = 32, 7 + vocab_size, rank = 32_000, 256 + + head = DSparkMarkovHead(vocab_size, rank).to(device).eval() + scheduler = DSparkScheduler(proposal_length).to(device) + + base_logits = torch.randn( + requests, + vocab_size, + device=device, + dtype=torch.float16, + ) + previous_tokens = torch.randint(vocab_size, (requests,), device=device) + confidence_logits = torch.randn( + requests, + proposal_length, + device=device, + dtype=torch.float16, + ) + + max_batch = requests * (proposal_length + 1) + batch_tokens = torch.arange(max_batch + 1, device=device) + profiled_steps_per_second = 1_000.0 / (1.0 + batch_tokens / 128.0) + + with torch.inference_mode(): + corrected_logits = head(base_logits, previous_tokens) + decision = scheduler(confidence_logits, profiled_steps_per_second) + + print("corrected logits:", corrected_logits.shape) + print("verification lengths:", decision.lengths) + print("expected tokens/step:", decision.expected_tokens) + print("expected tokens/second:", decision.expected_throughput) + + +if __name__ == "__main__": + main() diff --git a/Usage/DSpark/readme.md b/Usage/DSpark/readme.md new file mode 100644 index 0000000..a355c38 --- /dev/null +++ b/Usage/DSpark/readme.md @@ -0,0 +1,12 @@ +# DSpark usage + +Build the extension in `DSpark/`, then run from the repository root: + +```bash +python Usage/DSpark/dspark_usage.py +``` + +The example applies the fused low-rank Markov correction and feeds confidence +logits plus a target-engine step curve into the hardware-aware prefix scheduler. +See [`DSpark/README.md`](../../DSpark/README.md) for the full tensor contract, +weight import instructions, correctness tests, and benchmarks. diff --git a/pyproject.toml b/pyproject.toml index e3728bd..23eecf5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -31,6 +31,7 @@ dependencies = [ cuda = ["cupy-cuda12x>=12.0.0"] dev = ["pytest>=6.0", "pytest-cov", "black", "flake8"] examples = ["matplotlib>=3.0.0", "jupyter>=1.0.0"] +dspark = ["torch>=2.1.0"] [project.urls] Homepage = "https://github.com/dino65-dev/Cuda-ML-Library" @@ -38,4 +39,4 @@ Repository = "https://github.com/dino65-dev/Cuda-ML-Library.git" Issues = "https://github.com/dino65-dev/Cuda-ML-Library/issues" [tool.setuptools.dynamic] -version = {file = "VERSION"} \ No newline at end of file +version = {file = "VERSION"} diff --git a/setup.py b/setup.py index 82eca73..67641a8 100644 --- a/setup.py +++ b/setup.py @@ -42,7 +42,10 @@ def get_long_description(): 'matplotlib>=3.0.0', 'seaborn>=0.11.0', 'jupyter>=1.0.0', - ] + ], + 'dspark': [ + 'torch>=2.1.0', + ], } setup( @@ -63,6 +66,7 @@ def get_long_description(): package_data={ 'SVM': ['*.py'], 'HBM_SVM': ['*.py'], + 'DSpark': ['*.py', 'README.md'], 'Usage': ['**/*.py', '**/*.md'], }, include_package_data=True, @@ -96,4 +100,4 @@ def get_long_description(): ], }, zip_safe=False, # Due to shared libraries -) \ No newline at end of file +)