Skip to content

Sample the DSpark draft proposal from q, not argmax, for B2 #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
audreyt wants to merge 1 commit into
base: work-dspark
Choose a base branch
from
Open

Sample the DSpark draft proposal from q, not argmax, for B2 #1

Show file tree
Hide file tree
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
77 changes: 53 additions & 24 deletions ds4.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -10559,6 +10559,9 @@ static void output_logits_one_decode_scratch(

#ifndef DS4_NO_GPU
static int sample_argmax(const float *logits, uint32_t n_vocab);
int dspark_sample_draft_token(const float *logits, uint32_t vocab,
float temperature, uint64_t *rng,
float *scratch);

/* =========================================================================
* Metal Reference Comparison Helpers.
Expand Down Expand Up @@ -20329,7 +20332,9 @@ static bool metal_graph_eval_dspark_draft_block(
int *draft_n,
uint32_t *base_real_out,
float *last_logits,
float *all_draft_logits) {
float *all_draft_logits,
float draft_temperature,
uint64_t *draft_rng) {
if (draft_n) *draft_n = 0;
if (base_real_out) *base_real_out = 0;
if (!g || !target_model || !target_weights || !dspark_model || !mtp ||
Expand Down Expand Up @@ -20390,6 +20395,7 @@ static bool metal_graph_eval_dspark_draft_block(

const uint64_t row_bytes = (uint64_t)DS4_N_VOCAB * sizeof(float);
float *row_logits = xmalloc((size_t)row_bytes);
float *draft_scratch = draft_temperature > 0.0f ? xmalloc((size_t)row_bytes) : NULL;
for (uint32_t i = 0; ok && i < block_size; i++) {
ok = ds4_gpu_tensor_read(g->spec_logits,
(uint64_t)i * row_bytes,
Expand All @@ -20398,7 +20404,9 @@ static bool metal_graph_eval_dspark_draft_block(
if (!ok) break;
const int prev = i == 0 ? anchor_token : drafts[i - 1u];
dspark_apply_markov_bias(row_logits, dspark_model, mtp, prev);
drafts[i] = sample_argmax(row_logits, DS4_N_VOCAB);
drafts[i] = draft_temperature > 0.0f
? dspark_sample_draft_token(row_logits, DS4_N_VOCAB, draft_temperature, draft_rng, draft_scratch)
: sample_argmax(row_logits, DS4_N_VOCAB);
if (all_draft_logits) {
memcpy(all_draft_logits + (uint64_t)i * DS4_N_VOCAB, row_logits, (size_t)row_bytes);
}
Expand All @@ -20407,6 +20415,7 @@ static bool metal_graph_eval_dspark_draft_block(
}
}
free(row_logits);
free(draft_scratch);
if (!ok) return false;
*draft_n = (int)block_size;
return true;
Expand Down Expand Up @@ -23779,6 +23788,25 @@ static int b2_sample_from_log_probs(const float *log_probs, uint32_t vocab,
return (int)(vocab - 1);
}

/* Sample one token from the temperature-scaled softmax of `logits` (full
* vocab, no top-k/top-p truncation) using the xorshift64* state in `*rng`.
* `scratch` must point to a caller-owned buffer of `vocab` floats (reused
* across calls to avoid hot-path allocation); its contents are clobbered.
* Used to draw the DSpark draft proposal from the drafter distribution q
* when B2 mode is active, so the B2 accept/reject math below (which assumes
* proposals are actually samples from q, not argmax picks) is valid. */
int dspark_sample_draft_token(const float *logits, uint32_t vocab,
float temperature, uint64_t *rng,
float *scratch) {
if (temperature <= 0.0f) return sample_argmax(logits, vocab);
const float inv_temp = 1.0f / temperature;
for (uint32_t i = 0; i < vocab; i++) {
scratch[i] = logits[i] * inv_temp;
}
b2_log_softmax(scratch, vocab, scratch);
return b2_sample_from_log_probs(scratch, vocab, rng);
}

/* Sample from the residual distribution max(0, target_prob - draft_prob).
* Both inputs are log-probability vectors. */
static int b2_sample_residual(const float *log_target, const float *log_draft,
Expand Down Expand Up @@ -23836,7 +23864,7 @@ typedef struct {
* temperature: sampling temperature (<=0 falls back to argmax matching)
* rng: pointer to xorshift64* state (mutated)
*/
static b2_result b2_rejection_sample(
b2_result b2_rejection_sample(
const int *draft_tokens,
const float *draft_logits,
const float *target_logits,
Expand Down Expand Up @@ -24577,6 +24605,7 @@ struct ds4_session {
int dspark_draft_count;
uint32_t dspark_draft_base_real;
float *dspark_b2_draft_logits; /* [block_size * DS4_N_VOCAB] post-markov-bias logits for B2 */
float dspark_b2_temp; /* DS4_SPEC_TEMP resolved once at session create */
uint64_t dspark_b2_rng; /* xorshift64* state for B2 rejection sampling (persisted across calls) */
int dspark_prev_accepted; /* previous cycle accepted count (for adaptive block size) */
int dspark_prev_drafted; /* previous cycle drafted count */
Expand Down Expand Up @@ -28104,12 +28133,27 @@ int ds4_session_create(ds4_session **out, ds4_engine *e, int ctx_size) {
if (ds4_engine_has_mtp(e)) {
s->mtp_logits = xmalloc((size_t)DS4_N_VOCAB * sizeof(s->mtp_logits[0]));
s->mtp_draft_token = -1;
/* Allocate B2 draft logits buffer when DS4_SPEC_TEMP is set and DSpark is active. */
if (e->mtp_weights.kind == DS4_MTP_DRAFT_DSPARK && getenv("DS4_SPEC_TEMP")) {
s->dspark_b2_temp = 0.0f;
const char *spec_temp_env = getenv("DS4_SPEC_TEMP");
if (spec_temp_env && spec_temp_env[0]) {
char *end = NULL;
float v = strtof(spec_temp_env, &end);
if (end != spec_temp_env && v > 0.0f) s->dspark_b2_temp = v;
}
if (e->mtp_weights.kind == DS4_MTP_DRAFT_DSPARK && s->dspark_b2_temp > 0.0f) {
const uint32_t block_size = e->mtp_weights.dspark.block_size > 0
? e->mtp_weights.dspark.block_size : 16;
s->dspark_b2_draft_logits = xmalloc(
(size_t)block_size * DS4_N_VOCAB * sizeof(float));
const char *seed_env = getenv("DS4_SPEC_RNG_SEED");
if (seed_env && seed_env[0]) {
s->dspark_b2_rng = (uint64_t)strtoull(seed_env, NULL, 0);
}
if (s->dspark_b2_rng == 0) {
s->dspark_b2_rng = (uint64_t)time(NULL) ^
((uint64_t)getpid() << 32) ^
(uint64_t)clock();
}
}
}
if (e->distributed.role == DS4_DISTRIBUTED_COORDINATOR) {
Expand Down Expand Up @@ -29167,7 +29211,9 @@ static int ds4_session_eval_internal(ds4_session *s, int token, bool probe_mtp,
&draft_n,
&base_real,
getenv("DS4_MTP_FULL_LOGITS") ? s->mtp_logits : NULL,
s->dspark_b2_draft_logits)) {
s->dspark_b2_draft_logits,
s->dspark_b2_temp,
&s->dspark_b2_rng)) {
s->dspark_draft_count = draft_n;
s->dspark_draft_base_real = base_real;
s->mtp_draft_token = draft_n > 0 ? s->dspark_draft_tokens[0] : -1;
Expand Down Expand Up @@ -29305,24 +29351,7 @@ int ds4_session_eval_speculative_argmax(ds4_session *s, int first_token,
* ARDD adversarial review fix: RNG state persisted in session struct
* (not reseeded per call) to avoid correlated random sequences when
* multiple speculative eval calls happen within the same second. */
float b2_temp = 0.0f;
const char *spec_temp_env = getenv("DS4_SPEC_TEMP");
if (spec_temp_env && spec_temp_env[0]) {
char *end = NULL;
float v = strtof(spec_temp_env, &end);
if (end != spec_temp_env && v > 0.0f) b2_temp = v;
}
if (b2_temp > 0.0f && s->dspark_b2_rng == 0) {
const char *seed_env = getenv("DS4_SPEC_RNG_SEED");
if (seed_env && seed_env[0]) {
s->dspark_b2_rng = (uint64_t)strtoull(seed_env, NULL, 0);
}
if (s->dspark_b2_rng == 0) {
s->dspark_b2_rng = (uint64_t)time(NULL) ^
((uint64_t)getpid() << 32) ^
(uint64_t)clock();
}
}
const float b2_temp = s->dspark_b2_temp;

/* Greedy first-draft check (common to both greedy and B2 paths).
* At temp=0 this is exact; at temp>0 it is a fast pre-filter —
Expand Down
102 changes: 102 additions & 0 deletions tests/ds4_test.c
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2424,6 +2424,107 @@ static int test_run_dspark_target_cache_cli_missing_target_model(const char *dat
if (!WIFEXITED(status)) return -1;
return WEXITSTATUS(status);
}

/* Extern declarations for the internal (non-static) B2 helpers in ds4.c —
* intentionally not part of the public ds4.h surface, exposed just enough
* for this white-box statistical test. */
typedef struct {
int n_accepted;
int accepted_tokens[16];
int correction_token;
bool has_correction;
} b2_result;

int dspark_sample_draft_token(const float *logits, uint32_t vocab,
float temperature, uint64_t *rng,
float *scratch);

b2_result b2_rejection_sample(
const int *draft_tokens,
const float *draft_logits,
const float *target_logits,
uint32_t vocab,
int n_draft,
float temperature,
uint64_t *rng);

/* B2 rejection sampling must reproduce the target distribution exactly when
* the draft token is genuinely sampled from the proposal distribution q —
* that is the lossless-speculative-sampling invariant for n_draft=1
* (Chen et al. 2023 / Leviathan et al. 2023). This is a synthetic, CPU-only
* test: no GPU or model required. */
static void test_dspark_b2_rejection_sampling_unbiased(void) {
const uint32_t vocab = 6;
/* Deliberately different target (p) and draft (q) logits with full
* support, so the test actually exercises rejection + correction, not
* just the trivial all-accept case. */
const float target_logits[6] = { 2.0f, 0.5f, -1.0f, 1.0f, 0.0f, -0.5f };
const float draft_logits[6] = { 0.2f, 1.8f, 0.5f, -0.5f, 1.0f, 0.0f };
const float temperature = 1.0f;

/* Ground truth: softmax(target_logits / temperature). */
float target_probs[6];
{
float max_v = target_logits[0];
for (uint32_t i = 1; i < vocab; i++) if (target_logits[i] > max_v) max_v = target_logits[i];
float sum = 0.0f;
for (uint32_t i = 0; i < vocab; i++) {
target_probs[i] = expf((target_logits[i] - max_v) / temperature);
sum += target_probs[i];
}
for (uint32_t i = 0; i < vocab; i++) target_probs[i] /= sum;
}

const int N = 50000;
int hist_correct[6] = {0};
int hist_biased[6] = {0};
uint64_t rng_correct = 0xC0FFEEULL;
uint64_t rng_biased = 0xC0FFEEULL;
float scratch[6];

for (int trial = 0; trial < N; trial++) {
/* Correct path: proposal genuinely sampled from q. */
int draft_tok = dspark_sample_draft_token(draft_logits, vocab, temperature, &rng_correct, scratch);
b2_result r = b2_rejection_sample(
&draft_tok, draft_logits, target_logits, vocab, 1, temperature, &rng_correct);
int out_tok = r.has_correction ? r.correction_token : r.accepted_tokens[0];
TEST_ASSERT(out_tok >= 0 && out_tok < (int)vocab);
hist_correct[out_tok]++;

/* Biased control: proposal is argmax(q) (the bug this PR fixes), not
* a real sample. This MUST diverge from the target distribution,
* proving this test harness actually discriminates correct from
* broken proposal sampling. */
int argmax_tok = 0;
float best = draft_logits[0];
for (uint32_t i = 1; i < vocab; i++) if (draft_logits[i] > best) { best = draft_logits[i]; argmax_tok = (int)i; }
b2_result rb = b2_rejection_sample(
&argmax_tok, draft_logits, target_logits, vocab, 1, temperature, &rng_biased);
int out_biased = rb.has_correction ? rb.correction_token : rb.accepted_tokens[0];
TEST_ASSERT(out_biased >= 0 && out_biased < (int)vocab);
hist_biased[out_biased]++;
}

float max_dev_correct = 0.0f, max_dev_biased = 0.0f;
for (uint32_t i = 0; i < vocab; i++) {
float p_correct = (float)hist_correct[i] / (float)N;
float p_biased = (float)hist_biased[i] / (float)N;
float dev_correct = fabsf(p_correct - target_probs[i]);
float dev_biased = fabsf(p_biased - target_probs[i]);
if (dev_correct > max_dev_correct) max_dev_correct = dev_correct;
if (dev_biased > max_dev_biased) max_dev_biased = dev_biased;
}
fprintf(stderr,
"ds4-test: dspark-b2-unbiased max_dev_correct=%.4f max_dev_biased=%.4f (N=%d)\n",
max_dev_correct, max_dev_biased, N);
/* Correct (truly-sampled proposal) path must track the target
* distribution tightly: binomial std error at N=50000 is ~0.002-0.003,
* so 0.015 is a generous but still discriminating bound. */
TEST_ASSERT(max_dev_correct < 0.015f);
/* The argmax-proposal control must diverge well beyond that bound,
* proving the test can tell correct from broken proposal sampling. */
TEST_ASSERT(max_dev_biased > 0.05f);
}
static bool test_json_u64_field(const char *json, const char *key, uint64_t *out) {
const char *p = strstr(json, key);
if (!p) return false;
Expand Down Expand Up @@ -2582,6 +2683,7 @@ static const ds4_test_entry test_entries[] = {
{"--dspark-binder", "dspark-binder", "DSpark draft kind/config defaults without GGUF", test_dspark_binder_helpers},
{"--dspark-markov-bf16", "dspark-markov-bf16", "DSpark Markov BF16 tensor decoding", test_dspark_markov_bf16_helpers},
{"--dspark-runtime", "dspark-runtime", "DSpark capture plan and speculative gate helpers", test_dspark_runtime_helpers},
{"--dspark-b2-unbiased", "dspark-b2-unbiased", "DSpark B2 rejection sampling is unbiased vs target distribution", test_dspark_b2_rejection_sampling_unbiased},

{"--server", "server", "server parser/rendering/cache unit tests", test_server_unit_group},
};
Expand Down