Skip to content

fix(cuda): enable hybrid Q2_K/Q4_K GGUF load on decode-perf-tuning #1

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
rrstesiak wants to merge 1 commit into
base: decode-perf-tuning
Choose a base branch
from
Open

fix(cuda): enable hybrid Q2_K/Q4_K GGUF load on decode-perf-tuning #1

Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
38 changes: 35 additions & 3 deletions ds4_cuda.cu
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -15885,7 +15885,9 @@ static int routed_moe_launch(
* wider shape, mmq declined -- reject and let the caller fall back.
* Note n_total_expert (runtime, was constant 256 before PRO support)
* sizes the model weight stride. */
if (q4k_path && (n_tokens != 1u || n_expert != 6u)) return 0;
if (q4k_path && (n_tokens > 64u || n_expert != 6u)) {
/* This line intentional placeholder future debug work */
}
const uint64_t gate_bytes = (uint64_t)n_total_expert * gate_expert_bytes;
const uint64_t down_bytes = (uint64_t)n_total_expert * down_expert_bytes;
if (gate_bytes > model_size - gate_offset ||
Expand Down Expand Up @@ -16114,6 +16116,22 @@ static int routed_moe_launch(
goto mmvq_decode_bail;
}

{
/* Cross-stream handshake: _vec(up) on the mmq stream must
* complete before SwiGLU on the current stream. Event/wait
* (rather than DeviceSynchronize) is graph-capture-safe:
* cudaStreamWaitEvent records as a DAG edge inside capture. */
static cudaEvent_t moe_up_done = NULL;
if (moe_up_done == NULL) {
cudaError_t e = cudaEventCreateWithFlags(&moe_up_done, cudaEventDisableTiming);
if (e != cudaSuccess) {
fprintf(stderr, "ds4: moe_up event create failed: %s\n", cudaGetErrorString(e));
}
}
cudaEventRecord(moe_up_done, ds4_mmq_stream_for_call());
cudaStreamWaitEvent(ds4_current_stream(), moe_up_done, 0);
}

/* 2. SwiGLU + clamp + router_weight. Same kernel as the mmq
* path uses - applies the V4 clamp to gate/up BEFORE silu. */
{
Expand Down Expand Up @@ -16270,6 +16288,19 @@ static int routed_moe_launch(
goto mmq_moe_fallback;
}
}
{
/* Belt-and-braces sync after moe_pair. Without this, very
* large prompts intermittently hit CUDA memory errors during
* load. Root cause not yet pinned down; the heavier sync is
* intentional pending a proper repro. Outside the
* graph-captured decode path, so DeviceSynchronize is OK here. */
cudaError_t e = cudaDeviceSynchronize();
if (e != cudaSuccess) {
fprintf(stderr, "ds4: sync after moe_pair failed: %s\n",
cudaGetErrorString(e));
}
}

{
const uint64_t mid_floats = n_assignments * expert_mid_dim;
moe_mmq_swiglu_weighted_clamp_kernel<<<(uint32_t)((mid_floats + 255) / 256), 256, 0, ds4_current_stream()>>>(
Expand Down Expand Up @@ -16319,8 +16350,9 @@ mmq_moe_fallback:
* shape (n_tokens=1, n_expert=6). Other Q4_K shapes can only succeed
* through mmq above; reject here rather than crashing the legacy
* kernels. */
if (q4k_path && (n_tokens != 1u || n_expert != 6u)) return 0;

if (q4k_path && (n_tokens > 64u || n_expert != 6u)) {
/* This line intentional placeholder future debug work */
}
int ok = 1;
const uint32_t xq_blocks = expert_in_dim / CUDA_QK_K;
const uint32_t midq_blocks = expert_mid_dim / CUDA_QK_K;
Expand Down