From b6e799507621b5c0f96ceb4104365ddcaef559b3 Mon Sep 17 00:00:00 2001 From: tspeaks Date: Tue, 8 Sep 2026 00:10:17 -0500 Subject: [PATCH] fix(kernel): order the batch-memcpy probe against the current stream _probe allocates its destination with torch.zeros, which enqueues the fill on the current stream, then enqueues the verification copy on a fresh probe stream. Nothing joins the two. When the current stream has a backlog the copy completes first on the independent probe stream and the fill lands on top of it, so the probe reads back zeros and load_batch_memcpy raises "cudaMemcpyBatchAsync probe copied wrong bytes" on a GPU that supports the API. OffloadMoeCache catches that and falls back to full-layer copies, so --moe-prefill-hit-d2d silently does nothing whenever the current stream is busy as the probe runs -- in practice during prefill warmup, which is exactly when the flag is first exercised. A cold process hides the bug: the first torch.zeros pays a cudaMalloc and the first torch.cuda.Stream() populates the per-device stream pool, and each of those synchronizes the device, draining the backlog before the copy is enqueued. That is why the probe passes when run standalone and fails inside a warmed-up server. Join the probe stream to the current stream before the copy. Co-Authored-By: Claude Opus 5 --- python/freetoken/kernel/batch_memcpy.py | 6 +++ tests/kernels/test_batch_memcpy.py | 71 +++++++++++++++++++++++++ 2 files changed, 77 insertions(+) create mode 100644 tests/kernels/test_batch_memcpy.py diff --git a/python/freetoken/kernel/batch_memcpy.py b/python/freetoken/kernel/batch_memcpy.py index b39e5cde2..67e2bf987 100644 --- a/python/freetoken/kernel/batch_memcpy.py +++ b/python/freetoken/kernel/batch_memcpy.py @@ -28,6 +28,12 @@ def _probe(fn) -> None: src = torch.arange(16, dtype=torch.uint8).pin_memory() dst = torch.zeros(16, dtype=torch.uint8, device="cuda") stream = torch.cuda.Stream() + # torch.zeros filled dst on the CURRENT stream; the copy below runs on this one. + # Without the join the fill can land after the copy and clobber it, leaving the + # probe to read back zeros -- a false "copied wrong bytes" on hardware that + # supports the API. A cold process hides this (the first torch.zeros and + # torch.cuda.Stream() each synchronize the device); a warmed-up server does not. + stream.wait_stream(torch.cuda.current_stream()) fn( torch.tensor([dst.data_ptr()]), torch.tensor([src.data_ptr()]), diff --git a/tests/kernels/test_batch_memcpy.py b/tests/kernels/test_batch_memcpy.py new file mode 100644 index 000000000..bc20c13c3 --- /dev/null +++ b/tests/kernels/test_batch_memcpy.py @@ -0,0 +1,71 @@ +"""The cudaMemcpyBatchAsync load probe must be ordered against the current stream. + +_probe allocates its destination with torch.zeros (a fill enqueued on the CURRENT +stream) but enqueues the copy on a fresh probe stream. With nothing joining the two, +the fill can land *after* the copy and clobber it, so the probe reads back zeros and +load_batch_memcpy raises "probe copied wrong bytes" on a GPU that supports the API +perfectly well. It only reproduces once the caching allocator and the stream pool are +warm -- in a cold process torch.zeros and torch.cuda.Stream() each synchronize the +device and hide the race, which is why the probe passes standalone but fails inside a +server that is mid-warmup. +""" + +from __future__ import annotations + +import os + +import pytest +import torch + +CUDA = pytest.mark.skipif(not torch.cuda.is_available(), reason="needs CUDA") +JIT = pytest.mark.skipif( + os.getenv("FREETOKEN_DISABLE_JIT", "").strip().lower() in {"1", "true", "yes", "on"}, + reason="batch_memcpy has no AOT prebuild; needs runtime JIT", +) + + +def _cuda_at_least(major: int, minor: int) -> bool: + cuda = torch.version.cuda + if cuda is None: + return False + return tuple(int(x) for x in cuda.split(".")[:2]) >= (major, minor) + + +BATCH_API = pytest.mark.skipif( + not _cuda_at_least(13, 0), reason="the cudaMemcpyBatchAsync binding needs CUDA >= 13.0" +) + +# ~1s of spin on the current stream. The probe's own copy is microseconds, so any +# ordering bug shows up as the whole backlog landing on top of the copied bytes. +_BACKLOG_CYCLES = 1_500_000_000 + + +def _warm_device_state() -> None: + """Reach the state a server is in at prefill warmup: the caching allocator holds a + free block of the probe's 16-byte size and the per-device stream pool is populated, + so neither torch.zeros nor torch.cuda.Stream() synchronizes on the probe's behalf.""" + for _ in range(4): + block = torch.zeros(16, dtype=torch.uint8, device="cuda") + del block + pinned = torch.arange(16, dtype=torch.uint8).pin_memory() + del pinned + pool = [torch.cuda.Stream() for _ in range(40)] + torch.cuda.synchronize() + del pool + + +@CUDA +@JIT +@BATCH_API +def test_probe_survives_a_busy_current_stream(): + from freetoken.kernel.batch_memcpy import load_batch_memcpy + + _warm_device_state() + try: + torch.cuda._sleep(_BACKLOG_CYCLES) + # Guard against a vacuous test: if something drained the queue the race is + # not being exercised at all and a pass would mean nothing. + assert not torch.cuda.current_stream().query(), "backlog drained before the probe ran" + load_batch_memcpy() + finally: + torch.cuda.synchronize()