Add SIP-level tensor parallelism, component registry YAML, VA offset verification

- DPPolicy: 3-level (sip/cube/pe), unified naming (column_wise/row_wise)
- PE_CPU: auto num_programs from cube shard count
- context.launch(): per-SIP KernelLaunchMsg with local va_base + auto local shape
- deploy_tensor: removed mmus param, MMU mapping is context-only responsibility
- ComponentRegistry: YAML-based lazy loading (components.yaml), impls→builtin rename
- VA offset bench + tests: 2D/1D, standard Triton kernel pattern

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>
This commit is contained in:
2026-03-26 01:13:17 -07:00
parent 08812eda58
commit 63669f82cb
35 changed files with 813 additions and 219 deletions
+54 -4
View File
@@ -140,16 +140,65 @@ class ComponentRegistry:
Resolution order for ComponentRegistry.create(node, overrides, ctx):
1. overrides[node.impl] — caller-injected override
2. _registry[node.impl] — globally registered impl
2. _registry[node.impl] — globally registered impl (lazy import)
3. Error — no fallback; every node must have an impl
Registry is populated from components.yaml via load_components_yaml().
Manual register() is still supported for tests and overrides.
"""
_registry: dict[str, type[ComponentBase]] = {}
_lazy: dict[str, str] = {} # impl → "module.path:ClassName"
_loaded: bool = False
@classmethod
def register(cls, impl: str, component_cls: type[ComponentBase]) -> None:
cls._registry[impl] = component_cls
@classmethod
def load_components_yaml(cls, path: str | None = None) -> None:
"""Load impl→class mappings from components.yaml. Lazy imports on first use."""
if cls._loaded:
return
import yaml
from pathlib import Path
if path is None:
# Search: project root (cwd), then relative to this file
candidates = [
Path.cwd() / "components.yaml",
Path(__file__).parent.parent.parent.parent / "components.yaml",
]
for p in candidates:
if p.exists():
path = str(p)
break
if path is None:
return
with open(path) as f:
spec = yaml.safe_load(f)
for impl, class_path in (spec.get("components") or {}).items():
cls._lazy[impl] = class_path
cls._loaded = True
@classmethod
def _resolve(cls, impl: str) -> type[ComponentBase] | None:
"""Resolve impl name: check _registry first, then lazy import from _lazy."""
if impl in cls._registry:
return cls._registry[impl]
if not cls._loaded:
cls.load_components_yaml()
class_path = cls._lazy.get(impl)
if class_path is None:
return None
import importlib
module_path, class_name = class_path.rsplit(":", 1)
mod = importlib.import_module(module_path)
component_cls = getattr(mod, class_name)
cls._registry[impl] = component_cls # cache for next lookup
return component_cls
@classmethod
def create(
cls,
@@ -159,9 +208,10 @@ class ComponentRegistry:
) -> ComponentBase:
if overrides and node.impl in overrides:
return overrides[node.impl](node, ctx)
if node.impl in cls._registry:
return cls._registry[node.impl](node, ctx)
component_cls = cls._resolve(node.impl)
if component_cls is not None:
return component_cls(node, ctx)
raise ValueError(
f"No component registered for impl '{node.impl}' (node: {node.id}). "
f"Register it in kernbench.components.impls.__init__."
f"Add it to components.yaml or call ComponentRegistry.register()."
)