Skip to content

feat(cuda): GPU-native fused filtered sampling & lazy host KV materialization - #673

Open
karan68 wants to merge 14 commits into
timtoole02:mainfrom
karan68:feat/cuda-gpu-sampling-and-lazy-kv
Open

feat(cuda): GPU-native fused filtered sampling & lazy host KV materialization#673
karan68 wants to merge 14 commits into
timtoole02:mainfrom
karan68:feat/cuda-gpu-sampling-and-lazy-kv

Conversation

@karan68

@karan68 karan68 commented Aug 18, 2026

Copy link
Copy Markdown
Collaborator

Summary

This PR resolves a major throughput bottleneck during non-greedy sampling on CUDA and eliminates redundant synchronous host memory copies during prompt prefill:

  1. GPU-Native Fused Filtered Sampling:

    • Implemented in providing parallel on-device logit filtering (, , ) with Gumbel noise perturbation and reduction.
    • Connected directly to the decode forward pipeline in so non-greedy sampling ( > 0, ext{min_p} > 0$) executes 100% on the GPU with zero device-to-host logit transfers (saving ~512 KB PCIe round-trip per token on large vocabularies).
    • Added unit test in .
  2. Lazy Host KV Materialization on GPU Prefill:

    • Removed eager synchronous D2H KV cache copies from , relying on and on-demand if a CPU fallback or prompt cache export occurs.
  3. Platform Repack Safety:

    • Guarded in with to ensure Linux x86 hosts do not attempt ARM repacking.
    • Ensured execution planner explicitly disables CPU repack in when CUDA resident decode is active.

Verification & Benchmarks

  • Test Suite: Ran on an NVIDIA L4 GPU (), verifying and .
  • Live Serving: Verified with TinyLlama Q8_0 on NVIDIA L4 GPU handling live streaming chat completions () with temperature and min-p sampling.

karan68 added 14 commits August 18, 2026 19:30
…ent CUDA decode

- Add to_q8_0_bytes on Q8_0PackedRows4 to allow bit-exact reconstruction of GGUF Q8_0 bytes on demand without retaining duplicate raw blocks in host RAM.
- Add CpuTensor::q8_0_raw_bytes and update resident CUDA engine raw extraction to return Cow<[u8]>.
- Preserve q8_0_blocks.is_none() invariant when PackedRows4 is loaded, ensuring all tensor_store memory tests pass.
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