[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:
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:
|
||||
|
||||
Reference in new issue
Block a user