fix(kernel): order the batch-memcpy probe against the current stream

[tspeaks]

_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.

[gdevenyi]

Reproduced, and it is biting a production deployment — so this is worth more than a tidy-up.

Box: 2 x RTX 6000 Ada (sm_89), TP=2, Qwen3.8-Flash-Next NVFP4, offload MoE backend, torch 2.11.0+cu130, CUDA 13. Server runs with --moe-prefill-hit-d2d and --prefill-warmup (the default), which is precisely the combination your analysis predicts will lose.

I ran _probe's body as shipped and with your stream.wait_stream(torch.cuda.current_stream()), varying how much work is already queued on the current stream (a loop of 2048x2048 matmuls):

backlog on the current stream	as shipped	with the join
0	PASS	PASS
1	PASS	PASS
8	FAIL	PASS
40	FAIL	PASS
150	FAIL	PASS

Eight queued matmuls is enough. Your explanation of why it looks fine standalone matches exactly: at backlog 0 and 1 it passes here too, because the first torch.zeros and the first torch.cuda.Stream() each synchronize the device.

Why it matters beyond the probe. OffloadMoeCache catches the exception and falls back to full-layer copies, and it does so silently — there is no log line, so an operator has no way to know the flag stopped doing anything. On this deployment --moe-prefill-hit-d2d is worth TTFT 2.8 s -> 2.1 s on a 1k prompt, and the batch-memcpy entry point is resolved lazily on the first prefill that has a hit, i.e. during or just after the 16.6 s prefill warmup — the busiest the stream ever is. I have no probe-failure line in months of production logs precisely because the fallback does not emit one.

Two suggestions, neither blocking:

1. Log the fallback. Even at info_rank0, one line saying the batch path was declined and why would have made this self-diagnosing. Right now the only symptom is TTFT quietly regressing to the non-d2d number.
2. The same pattern is worth a grep elsewhere: any probe that allocates with a torch.* constructor on the current stream and then verifies on a private stream has this race. This one is the instance that costs a flag.

Taking the one-line fix onto our deploy branch now.

🤖 Generated with Claude Code

https://claude.ai/code/session_0173pf9k9fSVtwbm3f898HDt