Optimize DSpark hot paths for T4

DSpark/.gitignore

@@ -0,0 +1,6 @@
+*.egg-info/
+.pytest_cache/
+__pycache__/
+build/
+dist/
+*.so

DSpark/MANIFEST.in

@@ -0,0 +1,4 @@
+include README.md
+include install.sh
+recursive-include csrc *.cpp *.cu *.h
+recursive-include tests *.py

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.

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",
 ]

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__":

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,

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 <typename scalar_t>
+__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<uint64_t*>(shared_bytes);
+    auto* candidate_ids = reinterpret_cast<int32_t*>(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<uint64_t>(__float_as_uint(prefix));
+            const uint32_t position_tie_break =
+                0xffffffffu - static_cast<uint32_t>(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<float>(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<uint32_t>(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<torch::Tensor> 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<size_t>(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<scalar_t><<<
+                    1,
+                    padded_count,
+                    shared_bytes,
+                    stream>>>(
+                    confidence_logits.data_ptr<scalar_t>(),
+                    step_curve.data_ptr<float>(),
+                    temperatures.data_ptr<float>(),
+                    lengths.data_ptr<int32_t>(),
+                    survival.data_ptr<float>(),
+                    selected_count.data_ptr<int32_t>(),
+                    expected_tokens.data_ptr<float>(),
+                    expected_throughput.data_ptr<float>(),
+                    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);

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<torch::Tensor> dspark_schedule_cuda(
     const torch::Tensor& confidence_logits,
     const torch::Tensor& step_curve,