[executorch][muse] Move off-graph KV to the batching path and test batching in OSS CI

[Gasoonjia]

Stack from ghstack (oldest at bottom):

• -> [executorch][muse] Move off-graph KV to the batching path and test batching in OSS CI #23375
• [executorch][muse] Add run_solo_batching #23374
• [executorch][muse] Add export_solo_batching.py #23373
• [executorch][cuda] Let CudaExecutor honor a program's prefill lower bound #23372
• [executorch][llm] Share the batching executors' session logic under extension/llm/batching/util #23371
• [executorch][cuda] Add CudaExecutor #23370
• [executorch][cuda] Add step_plan for the CUDA batching executor #23369
• [executorch][llm] Add Buck targets for extension/llm/batching #23368
• [executorch][cuda] Add CudaCellCache #23367
• [executorch][cuda] Lower off-graph KV in a cell layout #23366
• [executorch][cuda] Extract CudaKVPool from CudaSequenceKVCache #23365

The off-graph KV cache now reaches Muse Glimmer through the batching path
only (export_solo_batching.py + run_solo_batching). This removes the earlier
single-sequence wiring from everywhere else in MG and switches the OSS CI mode
to the batching path, which now also tests batching itself.

Removed, so these match their state before the off-graph stack:

• runtime/engine/muse_glimmer_engine.{h,cpp}, runtime/runners/solo.cpp and
  muse_glimmer_worker.cpp: the off-graph plan, cache building and install,
  reset handling, offgraph_initial_capacity, and the engine's part of CUDA
  graph capture with the off-graph cache;
• export/export_solo.py: --use-offgraph-kv-cache and its branches. One
  piece stays: _solo_constant_methods omits get_mutable_buffer_metadata
  when there is none, because export_solo_batching.py shares it for a model
  whose cache lives off-graph.

Kept as they are, because the batching path uses them: the off-graph source
transformation in source_transformations/cuda.py (with its tests in
test_cuda_pipeline.py), export_solo_batching.py and run_solo_batching.

The CUDA backend side of the stack (lowering, CudaKVPool, CudaSequenceKVCache,
CudaCellCache, delegate-driven steps, graph recapture) is unchanged; the
batching path uses it.

CI: solo-text-offgraph becomes solo-text-batching. It exports with
export_solo_batching.py and runs run_solo_batching three times:

• one prompt with eager decode;
• the same prompt with the captured decode graph (these two cover what
  solo-text-offgraph checked, "Paris");
• five prompts in one batch: France, Japan, Italy, France again, and a prompt
  longer than a forward (Germany), which prefills in slices beside the others'
  decodes.

Each run writes a report, and .ci/scripts/check_muse_glimmer_batching_report.py
checks three things:

1. every generation finishes normally and contains its answer (Paris, Tokyo,
   Rome, Paris, Berlin), and eager and graph decode generate the same tokens.
   Tokens are not compared across batches: the France prompt prefilled alone
   and the one prefilled beside others run GEMMs of different shapes, and on
   the real model bf16 rounding steers greedy decoding differently (both
   still answer Paris);
2. the weights load once: loading costs at most 1.25x the weights, and does
   not vary by more than 256 MiB across runs reserving 4 or 5 sessions;
3. memory behaves like an off-graph KV cache: the pool holds the cells in use,
   grows geometrically (at most twice what was needed) instead of reserving
   every session's context, its bytes equal rows x bytes per cell, and
   generation adds no more than the pool plus a 3 GiB activation slack.

The export's mode check now accepts this mode (it listed only solo-text and
dflash-image, so the old mode could not have passed it).

Differential Revision: D123051488

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/23375

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 10 New Failures, 3 Unclassified Failures

As of commit 59464d8 with merge base c16dd0a ():

NEW FAILURES - The following jobs have failed:

• pull / unittest / linux / linux-job (gh)
  backends/apple/coreml/test/test_torch_ops.py::TestTorchOps::test_dequantize_affine_c8w_embedding_c8w_linear_ios16
• pull / unittest / macos / macos-job (gh)
  backends/apple/coreml/test/test_torch_ops.py::TestTorchOps::test_dequantize_affine_c8w_embedding_c8w_linear_ios16
• pull / unittest / windows / windows-job (gh)
  .ci/scripts/tests/test_wheel_test_modules.py::TestDroppedTestModules::test_ci_entry_points_still_match_the_workflows
• pull / unittest-editable / linux / linux-job (gh)
  backends/apple/coreml/test/test_torch_ops.py::TestTorchOps::test_dequantize_affine_c8w_embedding_c8w_linear_ios16
• pull / unittest-editable / macos / macos-job (gh)
  backends/apple/coreml/test/test_torch_ops.py::TestTorchOps::test_dequantize_affine_c8w_embedding_c8w_linear_ios16
• pull / unittest-editable / windows / windows-job (gh)
  .ci/scripts/tests/test_wheel_test_modules.py::TestDroppedTestModules::test_ci_entry_points_still_match_the_workflows
• trunk / test-samsung-models-linux / linux-job (gh)
  test_inception_v4_fp16
• trunk / unittest-release / linux / linux-job (gh)
  backends/apple/coreml/test/test_torch_ops.py::TestTorchOps::test_dequantize_affine_c8w_embedding_c8w_linear_ios16
• trunk / unittest-release / macos / macos-job (gh)
  backends/apple/coreml/test/test_torch_ops.py::TestTorchOps::test_dequantize_affine_c8w_embedding_c8w_linear_ios16
• trunk / unittest-release / windows / windows-job (gh)
  .ci/scripts/tests/test_wheel_test_modules.py::TestDroppedTestModules::test_ci_entry_points_still_match_the_workflows

UNCLASSIFIED FAILURES - DrCI could not classify the following jobs because the workflow did not run on the merge base. The failures may be pre-existing on trunk or introduced by this PR:

• Test CUDA Builds / test-muse-glimmer-cuda-e2e-kquant-17gb-solo-text-batching / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  RuntimeError: Command bash /__w/_temp/exec_script failed with exit code 1
• Test CUDA Builds / unittest-cuda / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  backends/cuda/tests/test_offgraph_kv.py::OffGraphKVCellCompileTest::test_decode_and_prefill_from_two_tokens_lower
• Test CUDA Builds / unittest-cuda-runtime / linux-job (gh) (this job did not run on the merge base, so DrCI cannot tell whether the failure is pre-existing)
  [ FAILED ] CudaKVCacheTest.CacheOutlivesTheStreamItLastSteppedOn

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[github-actions]

This PR needs a release notes: label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.