[Perf][MOSS-TTS] Avoid synchronization in batched index transfers

Construct MTP row offsets and mixed prefill/decode stop indices on the CPU, then enqueue nonblocking device transfers before batched writes. Add CUDA regressions for synchronization, row placement, stop flags, and snapshot ownership.

Signed-off-by: Sy03 <1370724210@qq.com>
This commit is contained in:
Sy03 committed 2026-09-08 00:39:50 +08:00
1 parent 251ef4f48f
commit dcd6588a36
3 files changed
+60 -2

No files matched your search

@@ -1704,7 +1704,10 @@ class MossTTSLocalTalkerForGeneration(nn.Module):
self._batch_should_continue = active
else:
self._batch_should_continue = torch.ones(len(info_dicts), device=hidden.device, dtype=torch.bool)
indices = torch.tensor(should_indices, device=hidden.device, dtype=torch.long)
# Keep mixed prefill/decode output packing asynchronous.
indices = torch.tensor(should_indices, device="cpu", dtype=torch.long).to(
hidden.device, non_blocking=True
)
self._batch_should_continue.index_copy_(0, indices, active)
if not have_codes: