Skip to content

Fix stale/NaN output when multiple DML sessions share a data transfer - #71

Open
KenLagos wants to merge 1 commit into
mainfrom
fence_fix
Open

Fix stale/NaN output when multiple DML sessions share a data transfer#71
KenLagos wants to merge 1 commit into
mainfrom
fence_fix

Conversation

@KenLagos

Copy link
Copy Markdown
Collaborator

The factory-owned DMLDataTransfer is shared across all sessions, but it cached a single execution provider and routed every CopyTensors call through that one provider's ExecutionContext/queue/fence. Each session (e.g. a pipeline of separate text-encoder/UNet/VAE models) has its own context and fence, so a copy for one session could run on another session's context. Its readback then waited on the wrong queue's fence — one already signaled — so WaitForSignal returned without blocking and the copy read back GPU memory before the owning session's compute had produced it, yielding stale or uninitialized (NaN) results on later runs.

Route each copy to the execution provider that owns the tensor instead of a cached one. DMLDataTransfer now keeps a registry of live providers and resolves the owner per copy by decoding the tensor's allocation handle to its bucketized allocator and matching that allocator to a registered EP. Copies with no GPU tensor (or an unmatched owner) fall back to the attached provider. This ensures each readback waits on the fence that actually produced the data.

  • dml_data_transfer.{h,cc}: provider registry (Register/UnregisterProvider), ResolveOwningProvider, and per-copy routing in CopyTensorsImpl
  • dml_execution_provider.h: GetBucketizedAllocator accessor for owner matching
  • dml_factory.cc: register/unregister each EP in Create/ReleaseEpImpl

The factory-owned DMLDataTransfer is shared across all sessions, but it
cached a single execution provider and routed every CopyTensors call
through that one provider's ExecutionContext/queue/fence. Each session
(e.g. a pipeline of separate text-encoder/UNet/VAE models) has its own
context and fence, so a copy for one session could run on another
session's context. Its readback then waited on the wrong queue's fence —
one already signaled — so WaitForSignal returned without blocking and the
copy read back GPU memory before the owning session's compute had
produced it, yielding stale or uninitialized (NaN) results on later runs.

Route each copy to the execution provider that owns the tensor instead of
a cached one. DMLDataTransfer now keeps a registry of live providers and
resolves the owner per copy by decoding the tensor's allocation handle to
its bucketized allocator and matching that allocator to a registered EP.
Copies with no GPU tensor (or an unmatched owner) fall back to the
attached provider. This ensures each readback waits on the fence that
actually produced the data.

- dml_data_transfer.{h,cc}: provider registry (Register/UnregisterProvider),
  ResolveOwningProvider, and per-copy routing in CopyTensorsImpl
- dml_execution_provider.h: GetBucketizedAllocator accessor for owner matching
- dml_factory.cc: register/unregister each EP in Create/ReleaseEpImpl
@KenLagos
KenLagos requested a review from apwojcik July 30, 2026 01:26
@KenLagos KenLagos self-assigned this Jul 30, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant