Description
When running bt_run in --mega mode, the generated output frequently consists of uninitialized or stale tokens.
During the prompt processing phase in run_prompt, mega_decode_launch correctly increments the device state counter *rt.d_step. However, when entering the main generation loop, the host code assumes that the newly generated tokens will be written starting from index 0 of the ring buffer. Since *d_step is never reset after processing the prompt, the kernel instead writes the first generated token to d_ring[prompt_length % kRingCap].
Details
Steps to Reproduce
Run bt_run with the --mega flag using any valid prompt. The generated sequence immediately produces nonsense tokens.
Suggested Fix
Explicitly reset the device's step counter immediately before entering the generation loop in mega_mode so that the device kernel and the host's expected read index remain perfectly synchronized.
// ... prompt processing finishes
int32_t first = 0;
CUDA_CHECK(cudaStreamSynchronize(rt.st));
CUDA_CHECK(cudaMemcpy(&first, rt.d_tok, 4, cudaMemcpyDeviceToHost));
std::printf(" %d", first);
n_done = 1;
// FIX: Reset d_step to 0 before the generation loop starts
CUDA_CHECK(cudaMemsetAsync(rt.d_step, 0, 4, rt.st));
CUDA_CHECK(cudaEventRecord(t0, rt.st));
// ... generation loop beginsImpact
The host ends up reading uninitialized memory (or stale prompt tokens) from the beginning of the ring buffer, while the actual generated tokens are silently written elsewhere. As a result, generation in --mega mode is completely broken, causing the model to output garbage immediately after the prompt.
Description
When running
bt_runin--megamode, the generated output frequently consists of uninitialized or stale tokens.During the prompt processing phase in
run_prompt,mega_decode_launchcorrectly increments the device state counter*rt.d_step. However, when entering the main generation loop, the host code assumes that the newly generated tokens will be written starting from index0of the ring buffer. Since*d_stepis never reset after processing the prompt, the kernel instead writes the first generated token tod_ring[prompt_length % kRingCap].Details
Steps to Reproduce
Run
bt_runwith the--megaflag using any valid prompt. The generated sequence immediately produces nonsense tokens.Suggested Fix
Explicitly reset the device's step counter immediately before entering the generation loop in
mega_modeso that the device kernel and the host's expected read index remain perfectly synchronized.Impact
The host ends up reading uninitialized memory (or stale prompt tokens) from the beginning of the ring buffer, while the actual generated tokens are silently written elsewhere. As a result, generation in
--megamode is completely broken, causing the model to output garbage immediately after the prompt.