feat(plan): surface backend selection in ContractionPlan

Add `ContractionPlan.backend` reporting which executor will run the
contraction — "nki" when strategy is matmul/bmm and neuronxcc is
importable, else "pytorch". Resolves Phase 1 acceptance criterion
that plan_contraction must report dispatch="nki" when shapes qualify.

- trntensor/plan.py: new field + _backend_for() helper
- tests/test_plan.py: test_backend_reports_nki_or_pytorch
- docs/api/plan.md: documents the new field

45 CPU tests pass.

Author: scttfrdmn

docs/api/plan.md

@@ -21,6 +21,9 @@ plan.transA, plan.transB  # whether to pre-transpose before matmul
 Dataclass with the fields above plus:
 
 - `subscripts` — the original subscript string
+- `backend` — executor that will run the contraction: `"nki"` (when
+  the strategy is `matmul`/`bmm` and `neuronxcc` is importable) or
+  `"pytorch"`
 - `estimated_flops: int`
 
 ## `estimate_flops(subscripts, *operands) -> int`

tests/test_plan.py

@@ -61,6 +61,23 @@ def test_batched_matmul_metadata(self):
         assert "j" in plan.contraction_indices
         assert "b" in plan.batch_indices
 
+    def test_backend_reports_nki_or_pytorch(self):
+        """`plan.backend` reports the executor that will run the contraction.
+
+        On Neuron hosts, matmul/bmm strategies report `"nki"`; on CPU they
+        report `"pytorch"`. The torch fallback always reports `"pytorch"`.
+        """
+        from trntensor.nki.dispatch import HAS_NKI
+
+        plan_mm = trntensor.plan_contraction(
+            "ij,jk->ik", torch.randn(4, 3), torch.randn(3, 5)
+        )
+        plan_torch = trntensor.plan_contraction(
+            "ij,jk,kl->il", torch.randn(3, 4), torch.randn(4, 5), torch.randn(5, 2)
+        )
+        assert plan_mm.backend == ("nki" if HAS_NKI else "pytorch")
+        assert plan_torch.backend == "pytorch"
+
 
 class TestEstimateFlops:
 

trntensor/plan.py

@@ -23,7 +23,8 @@
 class ContractionPlan:
     """Execution plan for a tensor contraction."""
     subscripts: str
-    strategy: str  # "matmul", "bmm", "torch", "nki"
+    strategy: str  # algorithm: "matmul" | "bmm" | "torch"
+    backend: str = "pytorch"  # executor: "nki" | "pytorch"
     transA: bool = False
     transB: bool = False
     contraction_indices: list[str] = field(default_factory=list)
@@ -32,15 +33,30 @@ class ContractionPlan:
     estimated_flops: int = 0
 
 
+def _backend_for(strategy: str) -> str:
+    """Resolve which executor will run a given strategy.
+
+    Returns ``"nki"`` only when the strategy maps to a NKI kernel
+    (``matmul`` or ``bmm``) and ``neuronxcc`` is importable.
+    """
+    if strategy in ("matmul", "bmm"):
+        from .nki.dispatch import HAS_NKI
+        if HAS_NKI:
+            return "nki"
+    return "pytorch"
+
+
 def plan_contraction(subscripts: str, *operands: torch.Tensor) -> ContractionPlan:
     """Analyze contraction and select execution strategy."""
     input_str, output_str = _parse_subscripts(subscripts)
     input_indices = input_str.split(",")
 
     if len(operands) == 2:
-        return _plan_binary(subscripts, input_indices, output_str, operands)
+        plan = _plan_binary(subscripts, input_indices, output_str, operands)
     else:
-        return ContractionPlan(subscripts=subscripts, strategy="torch")
+        plan = ContractionPlan(subscripts=subscripts, strategy="torch")
+    plan.backend = _backend_for(plan.strategy)
+    return plan
 
 
 def _parse_subscripts(subscripts: str) -> tuple[str, str]: