From fbebe2ec86172ee16923309f99c636a7b23dabe5 Mon Sep 17 00:00:00 2001 From: DINMAY KUMAR BRAHMA Date: Fri, 10 Jul 2026 21:58:53 +0530 Subject: [PATCH] Optimize DSpark T4 hot paths --- DSpark/.gitignore | 6 ++ DSpark/MANIFEST.in | 4 + DSpark/README.md | 48 +++++++---- DSpark/__init__.py | 2 + DSpark/benchmark.py | 21 +++-- DSpark/csrc/bindings.cpp | 54 +++++++++++- DSpark/csrc/dspark_cuda.cu | 158 ++++++++++++++++++++++++++++++++++++ DSpark/csrc/dspark_cuda.h | 6 ++ DSpark/dspark.py | 39 +++++++-- DSpark/install.sh | 2 +- DSpark/setup.py | 7 +- DSpark/tests/test_dspark.py | 43 ++++++---- README.md | 4 +- 13 files changed, 343 insertions(+), 51 deletions(-) create mode 100644 DSpark/.gitignore create mode 100644 DSpark/MANIFEST.in diff --git a/DSpark/.gitignore b/DSpark/.gitignore new file mode 100644 index 0000000..bf948b7 --- /dev/null +++ b/DSpark/.gitignore @@ -0,0 +1,6 @@ +*.egg-info/ +.pytest_cache/ +__pycache__/ +build/ +dist/ +*.so diff --git a/DSpark/MANIFEST.in b/DSpark/MANIFEST.in new file mode 100644 index 0000000..586cdc1 --- /dev/null +++ b/DSpark/MANIFEST.in @@ -0,0 +1,4 @@ +include README.md +include install.sh +recursive-include csrc *.cpp *.cu *.h +recursive-include tests *.py diff --git a/DSpark/README.md b/DSpark/README.md index 6cd5d30..7fbb6f9 100644 --- a/DSpark/README.md +++ b/DSpark/README.md @@ -27,12 +27,20 @@ For previous token `x`, the default head computes 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. +The optimized CUDA path stores `W2.T` as `[rank, vocab]`, gathers all `W1[x]` +rows, and submits one dense `addmm` update to cuBLAS: + +```text +output = 1 * (latent @ W2.T) + 1 * base_logits +``` + +This reuses projection tiles across requests, enables Tensor Cores for aligned +FP16/BF16 shapes, and folds the base-logit addition into the GEMM update. The +original scalar CUDA kernel remains available as `markov_logits_raw_cuda` for +research comparisons, but it is not the production default: treating a batch +as independent GEMVs rereads the projection matrix once per request and loses +Tensor Core throughput. During autograd the API uses native PyTorch operations +so parameter gradients remain exact. ### Confidence scheduler @@ -42,17 +50,21 @@ Given conditional acceptance estimates `c[r, k]`, calibrated prefix survival is p[r, k] = product(sigmoid(logit[r, j] / temperature[j]), j=0..k) ``` -The extension uses: +For up to 1024 candidates, the extension uses: + +- one fused CUDA launch for survival construction, sorting, admission, and + scatter; +- a shared-memory bitonic network with a 64-bit key whose low word enforces + prefix order on exact ties; and +- an exact causal first-throughput-drop scan matching DSpark Algorithm 1. -- 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. +Larger candidate sets automatically use the CUB radix-sort/scan implementation. +Both paths run on the current PyTorch stream without explicit 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. +For `R` requests and `K` draft positions, `K` may be 1–32; the paper's default +is 7. The fused path is selected when `R*K <= 1024`—including the common +`128*7=896` case. ## Install @@ -69,6 +81,10 @@ chmod +x install.sh ./install.sh ``` +Editable installation creates a `cuda_ml_dspark.egg-info/` directory containing +package metadata. It is normal, contains no runtime code or profiling data, and +is ignored by Git. + PyTorch chooses the local GPU architecture automatically. To build a portable binary, set `TORCH_CUDA_ARCH_LIST`, for example: @@ -163,6 +179,8 @@ cuBLAS depends on batch size, vocabulary, rank, dtype, and architecture. - 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. +- The optimized Markov path uses a dense cuBLAS/Tensor Core GEMM; the raw + scalar kernel is retained only for comparison. - 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. diff --git a/DSpark/__init__.py b/DSpark/__init__.py index 91ef924..9352f66 100644 --- a/DSpark/__init__.py +++ b/DSpark/__init__.py @@ -7,6 +7,7 @@ ScheduleResult, cuda_extension_available, markov_logits, + markov_logits_raw_cuda, schedule, ) @@ -17,5 +18,6 @@ "ScheduleResult", "cuda_extension_available", "markov_logits", + "markov_logits_raw_cuda", "schedule", ] diff --git a/DSpark/benchmark.py b/DSpark/benchmark.py index 063e64c..f0576cd 100644 --- a/DSpark/benchmark.py +++ b/DSpark/benchmark.py @@ -7,7 +7,12 @@ import torch -from .dspark import cuda_extension_available, markov_logits, schedule +from .dspark import ( + cuda_extension_available, + markov_logits, + markov_logits_raw_cuda, + schedule, +) from .reference import markov_logits_reference, schedule_reference @@ -89,6 +94,11 @@ def main() -> None: args.warmup, args.iterations, ) + raw_markov_us = _time_cuda( + lambda: markov_logits_raw_cuda(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, @@ -107,10 +117,11 @@ def main() -> None: 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") + print(f"Markov optimized: {custom_markov_us:9.2f} us") + print(f"Markov raw CUDA: {raw_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__": diff --git a/DSpark/csrc/bindings.cpp b/DSpark/csrc/bindings.cpp index 80136fa..bda9425 100644 --- a/DSpark/csrc/bindings.cpp +++ b/DSpark/csrc/bindings.cpp @@ -2,11 +2,63 @@ #include "dspark_cuda.h" +torch::Tensor dspark_markov_logits( + const torch::Tensor& base_logits, + const torch::Tensor& previous_token_ids, + const torch::Tensor& token_embedding, + const torch::Tensor& projection_t) { + TORCH_CHECK(base_logits.is_cuda(), "base_logits must be a CUDA tensor"); + TORCH_CHECK(previous_token_ids.is_cuda(), + "previous_token_ids must be a CUDA tensor"); + TORCH_CHECK(token_embedding.is_cuda(), + "token_embedding must be a CUDA tensor"); + TORCH_CHECK(projection_t.is_cuda(), "projection_t must be a CUDA tensor"); + 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(), + "Markov tensors 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(), + "Markov tensors must be on the same CUDA device"); + + const auto batch = base_logits.size(0); + const auto vocab = base_logits.size(1); + const auto rank = token_embedding.size(1); + TORCH_CHECK(previous_token_ids.size(0) == batch, + "previous_token_ids batch dimension mismatch"); + TORCH_CHECK(token_embedding.size(0) == vocab, + "token_embedding vocabulary mismatch"); + TORCH_CHECK(projection_t.size(0) == rank && + projection_t.size(1) == vocab, + "projection_t must have shape [rank, vocab]"); + + // One GEMM lets cuBLAS reuse W2 tiles across all requests and select a + // Tensor Core implementation. addmm folds base_logits into the GEMM + // update (beta=1), avoiding the separate vocabulary-sized add kernel used + // by the eager `base + latent @ projection` reference. + const auto latent = token_embedding.index_select(0, previous_token_ids); + return at::addmm(base_logits, latent, projection_t); +} + PYBIND11_MODULE(TORCH_EXTENSION_NAME, module) { module.def( "markov_logits", + &dspark_markov_logits, + "Tensor-Core DSpark low-rank Markov correction (CUDA)"); + module.def( + "markov_logits_raw", &dspark_markov_logits_cuda, - "Fused DSpark low-rank Markov logit correction (CUDA)"); + "Experimental scalar DSpark Markov correction (CUDA)"); module.def( "schedule", &dspark_schedule_cuda, diff --git a/DSpark/csrc/dspark_cuda.cu b/DSpark/csrc/dspark_cuda.cu index acfa4cd..fd5a295 100644 --- a/DSpark/csrc/dspark_cuda.cu +++ b/DSpark/csrc/dspark_cuda.cu @@ -145,6 +145,116 @@ __global__ __launch_bounds__(kThreads) void build_candidates_kernel( } } +// DSpark's production block is short (gamma=7), so common serving batches fit +// in one 1024-thread CTA. Keeping candidate construction, sorting, causal +// admission, and scatter in one launch avoids CUB temporary allocations and +// six extra kernel launches. Larger candidate sets use the CUB path below. +template +__global__ void schedule_small_kernel( + const scalar_t* __restrict__ logits, + const float* __restrict__ step_curve, + const float* __restrict__ temperatures, + int32_t* __restrict__ lengths, + float* __restrict__ survival, + int32_t* __restrict__ selected_count, + float* __restrict__ expected_tokens, + float* __restrict__ expected_throughput, + int request_count, + int proposal_length, + int candidate_count, + int padded_count) { + extern __shared__ unsigned char shared_bytes[]; + auto* keys = reinterpret_cast(shared_bytes); + auto* candidate_ids = reinterpret_cast(keys + padded_count); + const int tid = threadIdx.x; + + keys[tid] = 0; + candidate_ids[tid] = -1; + if (tid < request_count) { + lengths[tid] = 0; + } + __syncthreads(); + + // One thread builds the complete short prefix for one request. This avoids + // duplicated sigmoid work and is cheaper than remapping gamma-wide rows to + // warps when gamma is only seven. + for (int request = tid; request < request_count; request += blockDim.x) { + float prefix = 1.0f; + for (int position = 0; position < proposal_length; ++position) { + const int index = request * proposal_length + position; + const float scaled_logit = + as_float(logits[index]) / temperatures[position]; + prefix *= 1.0f / (1.0f + __expf(-scaled_logit)); + survival[index] = prefix; + + const uint64_t probability_bits = + static_cast(__float_as_uint(prefix)); + const uint32_t position_tie_break = + 0xffffffffu - static_cast(position); + keys[index] = (probability_bits << 32) | position_tie_break; + candidate_ids[index] = index; + } + } + __syncthreads(); + + // In-place ascending bitonic sort; thread zero consumes it in reverse. + // The composite key keeps earlier tokens ahead on equal probabilities. + for (int size = 2; size <= padded_count; size <<= 1) { + for (int stride = size >> 1; stride > 0; stride >>= 1) { + const int partner = tid ^ stride; + if (partner > tid) { + const bool ascending = (tid & size) == 0; + const uint64_t left = keys[tid]; + const uint64_t right = keys[partner]; + const bool swap = ascending ? (left > right) : (left < right); + if (swap) { + keys[tid] = right; + keys[partner] = left; + const int32_t left_id = candidate_ids[tid]; + candidate_ids[tid] = candidate_ids[partner]; + candidate_ids[partner] = left_id; + } + } + __syncthreads(); + } + } + + if (tid == 0) { + float expected = static_cast(request_count); + float best_throughput = expected * step_curve[request_count]; + int selected = 0; + + for (int order = 0; order < candidate_count; ++order) { + const int slot = padded_count - 1 - order; + const int candidate = candidate_ids[slot]; + const uint32_t probability_bits = + static_cast(keys[slot] >> 32); + const float probability = __uint_as_float(probability_bits); + if (candidate < 0 || !(probability > 0.0f)) { + break; + } + + const float candidate_expected = expected + probability; + const float candidate_throughput = + candidate_expected * step_curve[request_count + selected + 1]; + if (!(candidate_throughput > best_throughput)) { + break; + } + + expected = candidate_expected; + best_throughput = candidate_throughput; + ++selected; + const int request = candidate / proposal_length; + const int position = candidate - request * proposal_length; + lengths[request] = position + 1; + } + + selected_count[0] = selected; + expected_tokens[0] = expected; + expected_throughput[0] = best_throughput; + } +} + __global__ __launch_bounds__(kThreads) void unpack_probabilities_kernel( const uint64_t* __restrict__ sorted_keys, float* __restrict__ sorted_probabilities, @@ -360,6 +470,54 @@ std::vector dspark_schedule_cuda( const auto long_options = confidence_logits.options().dtype(torch::kInt64); const auto byte_options = confidence_logits.options().dtype(torch::kUInt8); + if (candidate_count <= 1024) { + int padded_count = 1; + while (padded_count < candidate_count) { + padded_count <<= 1; + } + auto survival = torch::empty_like(confidence_logits, float_options); + auto lengths = torch::empty({request_count}, int_options); + auto selected_count = torch::empty({1}, int_options); + auto expected_tokens = torch::empty({1}, float_options); + auto expected_throughput = torch::empty({1}, float_options); + const size_t shared_bytes = + static_cast(padded_count) * + (sizeof(uint64_t) + sizeof(int32_t)); + + AT_DISPATCH_FLOATING_TYPES_AND2( + torch::kFloat16, + torch::kBFloat16, + confidence_logits.scalar_type(), + "dspark_schedule_small_cuda", + [&] { + schedule_small_kernel<<< + 1, + padded_count, + shared_bytes, + stream>>>( + confidence_logits.data_ptr(), + step_curve.data_ptr(), + temperatures.data_ptr(), + lengths.data_ptr(), + survival.data_ptr(), + selected_count.data_ptr(), + expected_tokens.data_ptr(), + expected_throughput.data_ptr(), + request_count, + proposal_length, + candidate_count, + padded_count); + }); + C10_CUDA_KERNEL_LAUNCH_CHECK(); + return { + lengths, + survival, + selected_count, + expected_tokens, + expected_throughput, + }; + } + 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); diff --git a/DSpark/csrc/dspark_cuda.h b/DSpark/csrc/dspark_cuda.h index 13005c8..8b6d5f3 100644 --- a/DSpark/csrc/dspark_cuda.h +++ b/DSpark/csrc/dspark_cuda.h @@ -10,6 +10,12 @@ torch::Tensor dspark_markov_logits_cuda( const torch::Tensor& token_embedding, const torch::Tensor& projection_t); +torch::Tensor dspark_markov_logits( + 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, diff --git a/DSpark/dspark.py b/DSpark/dspark.py index 8a14547..aa7c7b8 100644 --- a/DSpark/dspark.py +++ b/DSpark/dspark.py @@ -109,6 +109,34 @@ def markov_logits( ) +def markov_logits_raw_cuda( + base_logits: torch.Tensor, + previous_token_ids: torch.Tensor, + token_embedding: torch.Tensor, + projection_t: torch.Tensor, +) -> torch.Tensor: + """Run the original scalar CUDA research kernel for comparison. + + This path intentionally remains available for microbatch experiments, but + it rereads the projection per request and is not the production default. + """ + + _check_markov_inputs( + base_logits, + previous_token_ids, + token_embedding, + projection_t, + ) + if not base_logits.is_cuda or _dspark_cuda is None: + raise RuntimeError("the compiled DSpark CUDA extension is required") + return _dspark_cuda.markov_logits_raw( + base_logits.contiguous(), + previous_token_ids.contiguous(), + token_embedding.contiguous(), + projection_t.contiguous(), + ) + + def _temperatures_tensor( temperatures: float | Sequence[float] | torch.Tensor | None, *, @@ -134,13 +162,10 @@ def _temperatures_tensor( ) 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()): + # Avoid launching a validation reduction on every decode step. Reusable + # DSparkScheduler temperatures are validated on CPU at construction time; + # direct CUDA callers are required to supply positive values. + if not result.is_cuda and not bool(torch.all(result > 0).item()): raise ValueError("all sequential calibration temperatures must be positive") return result.contiguous() diff --git a/DSpark/install.sh b/DSpark/install.sh index 70ebe9c..a55d726 100755 --- a/DSpark/install.sh +++ b/DSpark/install.sh @@ -2,4 +2,4 @@ set -euo pipefail cd "$(dirname "${BASH_SOURCE[0]}")" -python -m pip install --no-build-isolation -e . +python -m pip install --no-build-isolation --no-deps --force-reinstall -e . diff --git a/DSpark/setup.py b/DSpark/setup.py index 91d9dbc..52d472f 100644 --- a/DSpark/setup.py +++ b/DSpark/setup.py @@ -10,7 +10,7 @@ setup( name="cuda-ml-dspark", - version="0.1.0", + version="0.2.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", @@ -20,8 +20,8 @@ CUDAExtension( name="_dspark_cuda", sources=[ - str(ROOT / "csrc" / "bindings.cpp"), - str(ROOT / "csrc" / "dspark_cuda.cu"), + "csrc/bindings.cpp", + "csrc/dspark_cuda.cu", ], include_dirs=[str(ROOT / "csrc")], extra_compile_args={ @@ -38,6 +38,7 @@ ) ], cmdclass={"build_ext": BuildExtension.with_options(no_python_abi_suffix=True)}, + package_data={"DSpark": ["README.md"]}, 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 index 58d8036..fc92f8d 100644 --- a/DSpark/tests/test_dspark.py +++ b/DSpark/tests/test_dspark.py @@ -7,6 +7,7 @@ DSparkMarkovHead, cuda_extension_available, markov_logits, + markov_logits_raw_cuda, schedule, ) from DSpark.reference import markov_logits_reference, schedule_reference @@ -112,23 +113,31 @@ def test_cuda_matches_reference(dtype: torch.dtype) -> None: projection_t = torch.randn(16, 257, device=device, dtype=dtype) with torch.inference_mode(): actual = markov_logits(base, ids, embedding, projection_t) + raw = markov_logits_raw_cuda(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, - ) + torch.testing.assert_close(raw, expected, atol=tolerance, rtol=tolerance) + + # Exercise a small fused CTA, the exact 128x7 benchmark shape, and the + # >1024-candidate CUB fallback. + for requests in (17, 128, 160): + confidence = torch.randn(requests, 7, device=device, dtype=dtype) + temperatures = torch.ones(7, device=device) + curve = torch.ones(requests * 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/README.md b/README.md index 449e075..1772208 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ A high-performance CUDA-accelerated machine-learning kernel library with Python - **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 +- **DeepSeek DSpark**: Tensor-Core Markov logits and fused confidence-scheduled verification - **Memory Optimization**: Efficient GPU memory management with pooling - **Easy Integration**: Scikit-learn compatible API and PyTorch integration @@ -244,7 +244,7 @@ DSparkScheduler(proposal_length=7, temperatures=None) ``` - FP32, FP16, and BF16 inference kernels -- Warp-level calibrated prefix products and CUB candidate ranking +- Tensor-Core dense Markov update with fused single-CTA scheduling for common batches - 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