Out-of-bounds memory access (Buffer Overflow) in run_prompt for sequences > 1024 tokens

[kashifalikhan36]

Description

In tools/bt_run.cu's run_prompt(), the host streams tokens from the device ring buffer (rt.d_ring) using the launched variable to determine the copy size. When launched exceeds kRingCap (1024 tokens), cudaMemcpy attempts to read past the end of the rt.d_ring device allocation, which is strictly sized to 1024 tokens. Additionally, it writes past the bounds of the host-side ring vector.

Details

Steps to Reproduce

Run bt_run with the --mega or --graph flag and set the generation length to a value greater than 1024 (for example, --n 2000).

Suggested Fix

Bound the copy size to kRingCap and adjust the host-side array indexing to use modulo arithmetic, matching the kernel's circular buffer logic.

const int to_copy = std::min(launched, kRingCap);
CUDA_CHECK(cudaMemcpy(ring.data(), rt.d_ring, (size_t)to_copy * 4, cudaMemcpyDeviceToHost));

while (n_done - 1 < launched && n_done < n_gen) {
    const int32_t tok = ring[(size_t)((n_done - 1) % kRingCap)];
    // ...
}

Impact

This bug causes silent memory corruption on both the host and device, and can also result in a direct segmentation fault. As a result, it completely breaks long-context generation in mega and graph modes, making generation of sequences longer than 1024 tokens impossible.

[kashifalikhan36]

@jaberjaber23 Can you please assign me this issue, i can fix this ?