[executorch][cuda] Run packed noncausal DFlash draft steps on the cell cache

backends/cuda/CMakeLists.txt

@@ -509,5 +509,10 @@ if(BUILD_TESTING)
       aoti_cuda_backend
     )
     target_compile_definitions(test_cuda_kv_cache PRIVATE CUDA_AVAILABLE=1)
+    et_cxx_test(
+      test_cuda_cache_routing SOURCES runtime/test/test_cuda_cache_routing.cpp
+      EXTRA_LIBS aoti_cuda_backend
+    )
+    target_compile_definitions(test_cuda_cache_routing PRIVATE CUDA_AVAILABLE=1)
   endif()
 endif()

backends/cuda/cuda_backend.py

@@ -820,9 +820,18 @@ def get_custom_passes(cls, compile_specs: List[CompileSpec]) -> List[typing.Any]
                     )
                 triton_kernel_mode = mode
         passes = [MoveCondPredicateToCpuPass(), ReplaceInt64FloorDivWithFloatPass()]
+        cell_layout = any(
+            spec.key == OFFGRAPH_KV_COMPILE_SPEC
+            and parse_offgraph_kv_manifest(spec.value)["layout"] == "cell"
+            for spec in compile_specs
+        )
         for spec in compile_specs:
             if spec.key == OFFGRAPH_KV_STEP_WIDTH_COMPILE_SPEC:
-                passes.append(CheckOffGraphKVStepWidthPass(spec.value))
+                passes.append(
+                    CheckOffGraphKVStepWidthPass(
+                        spec.value, layout="cell" if cell_layout else "sequence"
+                    )
+                )
         for spec in compile_specs:
             if spec.key == OFFGRAPH_KV_COMPILE_SPEC:
                 passes.append(

backends/cuda/passes/lower_offgraph_kv.py

@@ -215,19 +215,20 @@ def cell_step(
     so no causal alignment describes visibility and the mask is always
     explicit. ``read_len`` bounds the sweep to the occupied extent.
 
-    ``cells`` and ``mask`` are declared at the widest step; this one uses
-    their first T rows. ``k``/``v`` are BHSD; the pools are BSHD. sdpa returns
+    ``cells`` and ``mask`` are declared at the widest step; writes use the
+    first K/V-width cells and attention uses the first query-width mask rows.
+    ``k``/``v`` are BHSD; the pools are BSHD. sdpa returns
     the query's dtype; ``out_dtype`` is the neutral op's requested output dtype.
     """
-    width = q.shape[2]
-    slots = cells[:width]
+    query_width = q.shape[2]
+    slots = cells[: k.shape[2]]
     _write(k_pool, slots, k)
     _write(v_pool, slots, v)
     out = torch.ops.triton.sdpa(
         q,
         k_pool.transpose(1, 2),
         v_pool.transpose(1, 2),
-        mask[:, :, :width, :],
+        mask[:, :, :query_width, :],
         0.0,
         False,
         scale,
@@ -319,6 +320,53 @@ def _source_input(node):
     return node
 
 
+def _cell_position_source(node):
+    # Edge conversion emits copy variants of the value-preserving view ops.
+    # Match overloads exactly rather than admitting arbitrary name substrings.
+    preserving = tuple(
+        op
+        for name in (
+            *_VALUE_PRESERVING_OPS,
+            "view_copy",
+            "unsqueeze_copy",
+            "squeeze_copy",
+            "expand_copy",
+            "alias_copy",
+        )
+        for op in (
+            getattr(torch.ops.aten, name).default,
+            getattr(exir_ops.edge.aten, name).default,
+        )
+    )
+    selects = (
+        torch.ops.aten.index_select.default,
+        exir_ops.edge.aten.index_select.default,
+    )
+    width = node.meta["val"].shape[0]
+    selected = False
+    while node.op == "call_function":
+        if node.target in preserving:
+            node = node.args[0]
+        elif node.target in selects and not selected:
+            source, dim, index = node.args
+            shape = source.meta["val"].shape
+            if (
+                dim != 0
+                or shape[0] != width
+                or tuple(node.meta["val"].shape) != tuple(shape)
+                or index.meta["val"].dim() != 1
+                or index.meta["val"].shape[0] != shape[0]
+            ):
+                raise ValueError(
+                    "off-graph cell position index_select must preserve full width on dim 0"
+                )
+            selected = True
+            node = source
+        else:
+            break
+    return node
+
+
 class CheckOffGraphKVStepWidthPass:
     """Check the declared step width against the delegated program.
 
@@ -331,7 +379,10 @@ class CheckOffGraphKVStepWidthPass:
 
     requires_exported_program = True
 
-    def __init__(self, value: bytes) -> None:
+    def __init__(self, value: bytes, layout: str = "sequence") -> None:
+        if layout not in ("sequence", "cell"):
+            raise ValueError(f"invalid off-graph KV layout {layout!r}")
+        self._cell_layout = layout == "cell"
         text = value.decode("utf-8")
         fields = text.split(":")
         if len(fields) != 2 or not all(f.isdigit() for f in fields):
@@ -358,11 +409,16 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram:
                 )
             declared = inputs[self._input]
             position = node.args[3]
-            if _source_input(position) is not declared:
+            source = (
+                _cell_position_source(position)
+                if self._cell_layout
+                else _source_input(position)
+            )
+            if source is not declared:
                 raise ValueError(
                     f"off-graph KV step width names input {self._input} "
                     f"({declared.name}), but the cache op's position comes "
-                    f"from {_source_input(position).name}"
+                    f"from {source.name}"
                 )
             shape = declared.meta["val"].shape
             width = position.meta["val"].shape[0]
@@ -372,6 +428,12 @@ def __call__(self, exported_program: ExportedProgram) -> ExportedProgram:
                     f"{declared.name} {tuple(shape)}, which is not the step's "
                     f"{width} positions"
                 )
+            k, v = (arg.meta["val"] for arg in node.args[1:3])
+            if k.shape[2] != width or v.shape[2] != width:
+                raise ValueError(
+                    "off-graph KV step width requires matching K/V extent and "
+                    f"positions, got {k.shape[2]}, {v.shape[2]} and {width}"
+                )
         return exported_program
 
 

backends/cuda/runtime/cuda_backend.cpp

@@ -109,6 +109,61 @@ Error commit_offgraph_kv_step(CudaDelegateHandle* handle, int64_t width) {
   return handle->kv_cache->commit_step(width);
 }
 
+// A per-method actual key overrides the artifact's logical role. Without that
+// override, the logical role names a load-time option containing the actual
+// key; it is never itself used as a process-global registry ID.
+Result<std::string> resolve_cuda_cache_key(
+    const BackendInitContext& context,
+    ArrayRef<CompileSpec> compile_specs) {
+  constexpr const char* option =
+      ::executorch::extension::llm::cache::kCacheKeyOption;
+  std::optional<std::string> logical;
+  for (const CompileSpec& spec : compile_specs) {
+    if (std::strcmp(spec.key, option) != 0) {
+      continue;
+    }
+    ET_CHECK_OR_RETURN_ERROR(
+        !logical && spec.value.buffer != nullptr && spec.value.nbytes > 0,
+        InvalidArgument,
+        "init: duplicate or empty logical CUDA cache key");
+    std::string_view value(
+        static_cast<const char*>(spec.value.buffer), spec.value.nbytes);
+    if (value.back() == '\0') {
+      value.remove_suffix(1);
+    }
+    ET_CHECK_OR_RETURN_ERROR(
+        !value.empty() && value.size() < runtime::kMaxOptionKeyLength &&
+            value.find('\0') == std::string_view::npos && value != option,
+        InvalidArgument,
+        "init: invalid logical CUDA cache key");
+    logical = std::string(value);
+  }
+  const auto actual = context.get_runtime_spec<const char*>(option);
+  ET_CHECK_OR_RETURN_ERROR(
+      actual.ok() || actual.error() == Error::NotFound || !logical,
+      InvalidArgument,
+      "init: CUDA cache runtime override must be a string");
+  if (actual.ok() && actual.get() != nullptr && *actual.get() != '\0') {
+    return std::string(actual.get());
+  }
+  if (!logical) {
+    return std::string{};
+  }
+  const auto mapped = context.get_runtime_spec<const char*>(logical->c_str());
+  ET_CHECK_OK_OR_RETURN_ERROR(
+      mapped.error(),
+      "init: logical CUDA cache '%s' requires a string runtime mapping to an "
+      "installed cache key, or a per-method '%s' override",
+      logical->c_str(),
+      option);
+  ET_CHECK_OR_RETURN_ERROR(
+      mapped.get() != nullptr && *mapped.get() != '\0',
+      InvalidArgument,
+      "init: logical CUDA cache '%s' has an empty runtime mapping",
+      logical->c_str());
+  return std::string(mapped.get());
+}
+
 class ET_EXPERIMENTAL CudaBackend final
     : public ::executorch::runtime::BackendInterface {
  private:
@@ -350,6 +405,10 @@ class ET_EXPERIMENTAL CudaBackend final
       }
     }
 
+    const auto resolved_cache_key =
+        resolve_cuda_cache_key(context, compile_specs);
+    ET_CHECK_OK_OR_RETURN_ERROR(resolved_cache_key.error());
+
     std::string so_blob_key;
     std::string weights_blob_key;
     CudaWeightCache::Metadata fqn_weights;
@@ -486,21 +545,16 @@ class ET_EXPERIMENTAL CudaBackend final
     // Runtime-owned off-graph buffers must capture their AOTI names before
     // the serialized constants update installs the ordinary weight set.
     //
-    // The cache is found by the key the runner published as a runtime spec.
-    // No key means an in-graph model, whose KV state is ordinary mutable
-    // buffers -- that is the whole of the off-graph/in-graph detection, and it
-    // needs nothing from the user.
+    // Resolve each delegate's logical cache role through this method's load
+    // options. The explicit actual-key override retains the ordinary runner
+    // contract; a logical name is never used directly as a registry ID.
     //
     // The key alone does not decide it, though: a program whose storage was
     // lowered off-graph carries no bytes for it, so running one without a cache
     // would hand its kernels unbound memory. That requirement is read from the
     // program's own constants and checked whether or not a key was passed.
-    const char* cache_key = nullptr;
-    if (auto spec = context.get_runtime_spec<const char*>(
-            ::executorch::extension::llm::cache::kCacheKeyOption);
-        spec.ok() && spec.get() != nullptr && *spec.get() != '\0') {
-      cache_key = spec.get();
-    }
+    const char* cache_key =
+        resolved_cache_key->empty() ? nullptr : resolved_cache_key->c_str();
     const auto needs_offgraph_kv = requires_offgraph_kv_storage(*handle);
     ET_CHECK_OK_OR_RETURN_ERROR(needs_offgraph_kv.error());
     if (cache_key != nullptr) {