Skip to content

Optimize DSpark hot paths for T4#8

Merged
dino65-dev merged 1 commit into
mainfrom
agent/optimize-dspark-t4
Jul 10, 2026
Merged

Optimize DSpark hot paths for T4#8
dino65-dev merged 1 commit into
mainfrom
agent/optimize-dspark-t4

Conversation

@dino65-dev

Copy link
Copy Markdown
Owner

Measured problem

Tesla T4 benchmark from the merged DSpark implementation:

Path CUDA PyTorch
Markov correction 10,756.48 us 292.56 us
Scheduler 254.77 us 1,191.04 us

The scheduler was already 4.7x faster, but the scalar Markov kernel was 36.8x slower.

Root cause

The raw Markov kernel decomposed a dense [128,256] @ [256,32000] operation into independent GEMVs. That reread the 16 MiB projection for each request and accumulated through scalar FP32 CUDA-core FMAs. PyTorch used one tiled GEMM with projection reuse and T4 Tensor Cores.

Changes

  • Make a single addmm/cuBLAS GEMM the production Markov path:
    • gathers all Markov embeddings once;
    • reuses projection tiles across requests;
    • permits Tensor Core selection for aligned FP16/BF16 dimensions;
    • folds the base-logit addition into the GEMM update.
  • Keep the old scalar kernel as explicit markov_logits_raw_cuda research baseline.
  • Add a one-launch scheduler path when requests * proposal_length <= 1024:
    • survival construction;
    • shared-memory bitonic candidate sort;
    • exact causal first-throughput-drop scan;
    • device-side result scatter.
  • Retain CUB as the large-candidate fallback.
  • Remove the per-call CUDA temperature-validation reduction.
  • Expand CUDA correctness coverage to 17x7, the measured 128x7 shape, and a 160x7 CUB fallback.
  • Fix editable/source packaging:
    • ignore generated *.egg-info/;
    • use relative extension source paths;
    • include the CUDA header in source distributions;
    • force extension rebuild in install.sh;
    • bump the DSpark package to 0.2.0.

Validation performed

  • python -m pytest DSpark/tests -q: 6 passed, 2 CUDA-only tests skipped locally
  • C++/PyBind syntax compilation against PyTorch headers: passed
  • Bitonic network tested against a reference sort through 1024 entries
  • Python bytecode and git diff --check: passed
  • Remote tree matches locally validated tree: 62fe74a6cc85d3badbc5e13a4d849c4ea25407b3

The CUDA-specific tests and benchmark must be rerun on the Tesla T4 before this draft is marked ready.

@dino65-dev dino65-dev marked this pull request as ready for review July 10, 2026 16:42
@dino65-dev dino65-dev merged commit b6aa61a into main Jul 10, 2026
@dino65-dev dino65-dev deleted the agent/optimize-dspark-t4 branch July 10, 2026 16:42
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant