[PLE] The pull registry refuses unbound or clashing names; test fixes
A schema walk fails on any pointer-like int parameter of a custom op on the PLE path. The FULL_K test child calls the Marlin MoE GEMM by keyword against its v0.30.0 signature. Co-Authored-By: Hinata <hinata@idihore.id> Signed-off-by: MOVZX <movzx@yahoo.com>
This commit is contained in:
@@ -428,7 +428,21 @@ class PleOffloadLayer(nn.Module, ABC):
|
||||
self._host_output_ptr = host_output_ptr
|
||||
self._host_row_bytes = int(host_row_bytes)
|
||||
if host_output_ptr is not None:
|
||||
_PULL_LAYERS[self._ple_offload_name] = self
|
||||
# Hardening: the registry key is the name the
|
||||
# traced graph carries, so it must be a real, unique layer name.
|
||||
name = self._ple_offload_name
|
||||
if name == "<unbound>":
|
||||
raise RuntimeError(
|
||||
"PLE pull: setup_cross_process_offload on a layer with no "
|
||||
"_ple_offload_name; the op would resolve nothing by name"
|
||||
)
|
||||
prior = _PULL_LAYERS.get(name)
|
||||
if prior is not None and prior is not self:
|
||||
raise RuntimeError(
|
||||
f"PLE pull: two layers armed under the name {name!r}; the op "
|
||||
"would resolve the wrong layer's pointers"
|
||||
)
|
||||
_PULL_LAYERS[name] = self
|
||||
|
||||
def forward(
|
||||
self,
|
||||
|
||||
Reference in New Issue
Block a user