diff --git a/.github/workflows/build-test.yml b/.github/workflows/build-test.yml index 7ffda68c..92d06fc6 100644 --- a/.github/workflows/build-test.yml +++ b/.github/workflows/build-test.yml @@ -21,6 +21,8 @@ jobs: steps: - name: Checkout Source Code uses: actions/checkout@v3 + with: + fetch-depth: 0 - name: Set up Python 3.10 uses: actions/setup-python@v5 @@ -30,7 +32,7 @@ jobs: - name: Install Dependencies run: | python -m pip install --upgrade pip - pip install "setuptools>=78.1.1,<82" wheel build + pip install "setuptools>=78.1.1,<82" "setuptools-scm>=8" wheel build pip install torch --index-url https://download.pytorch.org/whl/cpu - name: Build sdist diff --git a/.github/workflows/ci-pr.yml b/.github/workflows/ci-pr.yml index d40ab0f5..72d168ae 100644 --- a/.github/workflows/ci-pr.yml +++ b/.github/workflows/ci-pr.yml @@ -54,6 +54,8 @@ jobs: steps: - uses: actions/checkout@v4 + with: + fetch-depth: 0 - name: Set up Python ${{ matrix.python-version }} uses: actions/setup-python@v5 @@ -69,7 +71,7 @@ jobs: | pip install -r /dev/stdin 2>&1 || true - name: Install test tools - run: pip install pytest pytest-xdist pytest-timeout + run: pip install pytest pytest-xdist pytest-timeout "setuptools-scm>=8" - name: Install package (no CUDA extensions) run: pip install --no-build-isolation --no-deps -e . @@ -80,6 +82,11 @@ jobs: -m "not gpu and not multi_gpu and not slow and not perf and not network" \ -x --timeout=120 -v + - name: Run DFlash unified validation E2E + run: | + pytest tests/python/integration/test_dflash_unified_validation_e2e.py \ + -v + - name: Install contextpilot run: pip install contextpilot @@ -92,6 +99,8 @@ jobs: runs-on: ubuntu-22.04 steps: - uses: actions/checkout@v4 + with: + fetch-depth: 0 - name: Set up Python 3.10 uses: actions/setup-python@v5 @@ -100,7 +109,7 @@ jobs: - name: Install build tools run: | - pip install "setuptools>=78.1.1,<82" wheel build + pip install "setuptools>=78.1.1,<82" "setuptools-scm>=8" wheel build pip install torch --index-url https://download.pytorch.org/whl/cpu - name: Build sdist diff --git a/.github/workflows/publish-test.yml b/.github/workflows/publish-test.yml index 5b7091f5..36b7ea1f 100644 --- a/.github/workflows/publish-test.yml +++ b/.github/workflows/publish-test.yml @@ -13,36 +13,16 @@ on: - 'tests/**' - 'docs/**' -env: - # Bump this to the NEXT planned release when cutting a stable release. - # Nightlies must sort higher than the current stable so that - # `pip install --pre moe-infinity` picks them up. - NIGHTLY_BASE_VERSION: "0.0.2" - permissions: contents: write jobs: - setup-version: - runs-on: ubuntu-22.04 - steps: - - name: Generate version number - run: | - VERSION_HASH=$(date +"%Y%m%d%H%M%S") - echo "Generated version hash: $VERSION_HASH" - echo $VERSION_HASH > version.txt - - - name: Upload version number as artifact - uses: actions/upload-artifact@v4 - with: - name: version - path: version.txt - wheel: name: Build Wheel runs-on: ${{ matrix.os }} - needs: setup-version - permissions: write-all + permissions: + id-token: write + contents: read strategy: fail-fast: false matrix: @@ -53,12 +33,8 @@ jobs: steps: - name: Checkout Source Code uses: actions/checkout@v3 - - - name: Download version value artifact - uses: actions/download-artifact@v4 with: - name: version - path: artifact + fetch-depth: 0 - name: Free disk space run: | @@ -83,13 +59,16 @@ jobs: - name: Install Python build dependencies (torch matched to CUDA) run: | python3 -m pip install --upgrade pip - python3 -m pip install "setuptools>=78.1.1,<82" wheel ninja py-cpuinfo build + python3 -m pip install "setuptools>=78.1.1,<82" "setuptools-scm>=8" wheel ninja py-cpuinfo build python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu128 - name: Fetch CUTLASS headers run: | git clone --depth 1 --branch v3.9.2 https://github.com/NVIDIA/cutlass.git "${HOME}/cutlass" + # setuptools-scm derives the version from git history: a commit that is N + # commits after tag "vX.Y.Z" builds as "X.Y.(Z+1).devN". local_scheme in + # pyproject.toml strips the "+gHASH" suffix so the sdist is PyPI-valid. - name: Build Wheel shell: bash env: @@ -98,8 +77,7 @@ jobs: export CUDA_HOME="/usr/local/cuda-${{ matrix.cuda-version }}" export PATH="${CUDA_HOME}/bin:${PATH}" export CUTLASS_DIR="${HOME}/cutlass" - VERSION_HASH=$(cat artifact/version.txt) - MOEINF_VERSION=${{ env.NIGHTLY_BASE_VERSION }}.dev${VERSION_HASH} python3 -m build --wheel --no-isolation + python3 -m build --wheel --no-isolation wheel_name=$(ls dist/*whl | xargs -n 1 basename) asset_name=${wheel_name//"linux"/"manylinux1"} echo "wheel_name=${wheel_name}" >> $GITHUB_ENV @@ -108,8 +86,7 @@ jobs: - name: Build Source if: ${{ matrix.python-version == '3.10' }} run: | - VERSION_HASH=$(cat artifact/version.txt) - MOEINF_VERSION=${{ env.NIGHTLY_BASE_VERSION }}.dev${VERSION_HASH} python3 -m build --sdist + python3 -m build --sdist - name: Rename Wheel run: | @@ -122,12 +99,12 @@ jobs: cp dist/*.tar.gz sdist_only/ # Nightly publishes only the sdist to PyPI (see publish.yml for rationale). + # Auth via PyPI Trusted Publishing (OIDC) — no username/password needed; + # requires id-token: write (above) and a publisher registered on PyPI for + # this repo + workflow file. See RELEASE.md. - name: Publish sdist to PyPI if: ${{ github.event_name == 'push' && github.ref == 'refs/heads/main' && matrix.python-version == '3.10' }} uses: pypa/gh-action-pypi-publish@release/v1.8 with: skip-existing: true packages-dir: sdist_only - env: - TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} - TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} diff --git a/.github/workflows/publish.yml b/.github/workflows/publish.yml index 8f20ea04..34535e5e 100644 --- a/.github/workflows/publish.yml +++ b/.github/workflows/publish.yml @@ -7,6 +7,7 @@ on: permissions: contents: write + id-token: write jobs: release: @@ -47,6 +48,8 @@ jobs: steps: - name: Checkout Source Code uses: actions/checkout@v3 + with: + fetch-depth: 0 - name: Free Disk Space run: | @@ -71,7 +74,7 @@ jobs: - name: Install Python build dependencies (torch matched to CUDA) run: | python3 -m pip install --upgrade pip - python3 -m pip install "setuptools>=78.1.1,<82" wheel ninja py-cpuinfo build + python3 -m pip install "setuptools>=78.1.1,<82" "setuptools-scm>=8" wheel ninja py-cpuinfo build python3 -m pip install torch --index-url https://download.pytorch.org/whl/cu128 - name: Fetch CUTLASS headers @@ -120,12 +123,12 @@ jobs: # Only the sdist is published to PyPI. Prebuilt CUDA wheels are ABI-locked to a # specific torch/CUDA, which PyPI wheel tags cannot express; they are attached to # the GitHub Release instead. pip installs the sdist and compiles it locally. + # Auth via PyPI Trusted Publishing (OIDC) — no username/password needed; + # requires id-token: write (above) and a publisher registered on PyPI for + # this repo + workflow file. See RELEASE.md. - name: Publish sdist to PyPI if: ${{ matrix.python-version == '3.10' }} uses: pypa/gh-action-pypi-publish@release/v1.8 with: skip-existing: true packages-dir: sdist_only - env: - TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }} - TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }} diff --git a/.gitignore b/.gitignore index 214ead71..a8250e45 100644 --- a/.gitignore +++ b/.gitignore @@ -244,3 +244,6 @@ benchmarks/expert_io_microbench/results/ benchmarks/eval/__pycache__/ *.nsys-rep *.sqlite + +# setuptools-scm generated version file (regenerated on every build) +moe_infinity/_version.py diff --git a/ARCHITECTURE.md b/ARCHITECTURE.md index e2127ec4..33d15f30 100644 --- a/ARCHITECTURE.md +++ b/ARCHITECTURE.md @@ -47,7 +47,7 @@ moe_infinity/ │ │ block classes with Sync* wrappers, sets up │ │ expert tracing │ ├── attention_backend.py Attention backend dispatch (SDPA, FlashAttention, -│ │ FlashInfer, fallback backend) +│ │ FlashInfer where supported, fallback backend) │ ├── hooks.py Forward-pass hooks for tracing and prefetching │ └── compile.py Optional torch.jit compilation of expert MLPs │ @@ -74,6 +74,9 @@ moe_infinity/ │ ├── model_runner.py Runs a prefill/decode step for a batch │ ├── batch.py BatchBuilder + SchedulerOutput │ ├── kv_cache.py PagedKVCache, BlockAllocator, BlockTable +│ ├── spec_session_driver.py Persistent DFlash session lifecycle; Stage 4a/4b selection +│ ├── spec_cache_adapter.py Per-sequence adapter over engine-owned MLA pages +│ ├── mla_cache.py Packed DeepSeek latent/rope page storage │ ├── spec_state.py SpecDecodeState and committed-count bookkeeping │ ├── spec_verify.py apply_verify_step rollback helper │ ├── memory_manager.py GPU memory budget coordination @@ -89,9 +92,14 @@ moe_infinity/ │ ├── eviction_sync.py Request-termination → ContextPilot eviction │ └── contextpilot_*.py Optional prompt-optimization middleware │ -├── spec_decode/ DFlash speculative decoding and route-ahead support +├── spec_decode/ DFlash semantic core, execution backends, route-ahead │ ├── dflash.py DFlashSpeculator, config readers, validators, -│ │ route-ahead hook +│ │ canonical SpecSession transitions +│ ├── protocols.py Request, capability, trace, rich metadata contracts +│ ├── session_driver.py Capability selection and request cohorts +│ ├── backends.py Per-request execution protocol/backend +│ ├── backends_bare_hf.py Physical greedy/sampled/mixed dense cohorts +│ ├── backends_rich.py Capability-gated physical rich cohorts │ ├── _route_ahead_ctx.py Verify-time route-ahead contextvars and │ │ prefetch handle │ ├── _route_ahead_stats.py Opt-in route-ahead coverage/waste metrics @@ -136,10 +144,11 @@ moe_infinity/ ## 3. Two Execution Paths -MoE-Infinity has two runtime paths. The deprecated sync path currently owns -`MoE.generate(..., speculative_draft=...)`. The recommended continuous-batching -path is the async HTTP service started by `MoE.serve()` or `api_server_v2.py`; -it is not a drop-in in-process return API. +MoE-Infinity retains two scheduler/lifecycle paths, but DFlash no longer has two +acceptance/sampling implementations. Direct and deprecated-sync execution use +`SessionDriver`; serving drives the same canonical `SpecSession` transitions +through `SpecSessionDriver` while retaining scheduler and cancellation +ownership. `MoE.serve()` remains an HTTP service, not an in-process tensor API. ### Path A - Deprecated synchronous path (`MoE.generate()`) @@ -147,10 +156,14 @@ it is not a drop-in in-process return API. flowchart TD U[User code] --> M[MoE.generate(..., speculative_draft=...)] M --> R[_resolve_spec_strategy()] - R -->|None / False / non-greedy / batch>1| S[GenerationEngine._generate_standard()] + R -->|None / False / unsupported facade mode| S[GenerationEngine._generate_standard()] R -->|attach| E[GenerationEngine.spec_strategy] E --> D[DFlashSpeculator.generate()] - D --> F[MoE._native_model_forward_rich()] + D --> J[SessionDriver + SpecSession per row] + J --> B{backend capability} + B --> H[batched bare HF] + B --> P[grouped per-request rich] + B --> F[physical row-aware rich] F --> G[Sync* MoE blocks] G --> X[DistributedExpertExecutor.dispatch_local()] D --> K[accept, rollback, cache rewind] @@ -163,8 +176,9 @@ Notes: - `MoE.generate()` emits `DeprecationWarning` and is scheduled for removal. This diagram records the current transition path rather than a stable method contract. - `_resolve_spec_strategy()` attaches a speculator per call. Passing `None` or `False` detaches it and the standard path runs. -- The greedy gate lives in `GenerationEngine.generate()` and `_spec_strategy_applies()`. If the request is not a singleton greedy decode, the engine uses `_generate_standard()` and the output stays on the pre-DFlash baseline. -- `DFlashSpeculator._forward_target()` uses `moe._native_model_forward_rich()` when available, so expert dispatch and any configured prefetch hook stay intact. +- Direct bare-HF execution supports batch-1/batch>1 greedy, sampled, and mixed rows. The deprecated facade retains narrower compatibility gating and does not silently widen sampled batch > 1. +- `_generate_batched` is an adapter over `SessionDriver` and a physical backend, not an independent semantic loop. +- Rich rows are physically batched only when the wrapper declares the complete row-aware contract. MLA/hybrid wrappers otherwise use grouped per-request sessions. - The route-ahead context is verify-only. `DistributedExpertExecutor._maybe_route_ahead_prefetch()` reads the active context, computes the exact expert union from the router mask, and pins that exact set. If the context is inactive, no prefetcher is bound, or the union is empty, the legacy path is unchanged. - `SyncGptOssMLP` stays resident. When no executor seam exists, it only records read-only route-ahead stats. @@ -177,26 +191,34 @@ flowchart TD A --> S[Scheduler.schedule()] S --> B[BatchBuilder.from_scheduler_output()] B --> C{ContinuousBatchingEngine.step()\n_can_delegate_speculative(batch)?} - C -->|yes| D[ContinuousBatchingEngine._step_speculative()] + C -->|eligible| D[SpecSessionDriver persistent session] C -->|no| E[ContinuousBatchingEngine._execute_batch()] E --> N[Sampler.sample()] - D --> G[DFlashSpeculator.generate()] - D --> U[Scheduler.update_after_step(... committed_counts=...)] + D --> M{execution context} + M --> T[Stage 4a temporary_dynamic] + M --> P[Stage 4b paged_mla] + D --> U[publish canonical committed tokens] U --> Q[SequenceData + PagedKVCache] R --> X[StreamManager.push_token + SSE] X --> Y[abort_request() on disconnect or cleanup] ``` Notes: -- `Scheduler.schedule()` and `BatchBuilder.from_scheduler_output()` run before the speculative gate. `ContinuousBatchingEngine.step()` then checks `_can_delegate_speculative(batch)`; only eligible fresh singleton greedy prefill requests enter `_step_speculative()`, otherwise the normal `_execute_batch()` + sampler path runs. -- `_step_speculative()` emits one `RequestOutput` per committed token, then calls `scheduler.update_after_step(..., committed_counts={...})`. -- `Scheduler.update_after_step()` advances `SequenceData`, appends committed tokens through `PagedKVCache.append_tokens()`, and frees completed sequences with `PagedKVCache.free_sequence()`. +- Persistent sessions preserve row-local sampling policy and request RNG. Unsupported request metadata falls back before drafting; it is never silently converted to greedy. +- Stage 4a uses an explicit temporary private DynamicCache context. Stage 4b is default-off and restricted to eligible greedy batch-1 DeepSeek V2/V3 MLA. +- Task 8.5 established the DeepSeek MLA prerequisite before Stage 4b: the engine owns packed latent/rope target pages through `PagedCacheAdapter`; the draft cache remains separate. +- Standard and packed MLA target caches are distinct engine-owned paged stores. + Each eligible request has exactly one target store, while the drafter cache is + separate. All in-flight DRAFT/VERIFY sessions are resident and + non-preemptible; there is no speculative swap/resume claim. Cancellation + releases resources after an in-flight backend call returns. - `serving/spec_state.py` and `serving/spec_verify.py` define low-level tested helper contracts. `SpecDecodeState.record_verify()` and `apply_verify_step()` model committed-count bookkeeping for the rollback tests; the live `ContinuousBatchingEngine` path uses `Scheduler.update_after_step(..., committed_counts=...)` together with `PagedKVCache` directly. - Streaming and cleanup are ordinary server behavior. `StreamManager.push_token()` emits SSE chunks, and `_completion_event_generator()` / `_chat_event_generator()` call `abort_request()` when the client disconnects. ### Shared Components -Both paths share: +Both paths share the canonical DFlash protocol/trace types and session semantics, +as well as: - `runtime/model_offload.py` for model loading and MoE block monkey-patching - `runtime/attention_backend.py` for attention kernel dispatch - `memory/` for expert cache / KV cache bookkeeping @@ -210,9 +232,9 @@ Both paths share: 1. **Intake.** `api_server_v2` validates the request, tokenizes the prompt, and calls `engine.add_request(...)`, which creates a `SequenceData`. 2. **Scheduling.** On each async tick, `Scheduler.schedule()` decides which sequences to prefill, which to decode, and which to preempt. It allocates paged KV blocks through `PagedKVCache`. 3. **Batching.** `BatchBuilder.from_scheduler_output()` assembles packed input tensors, attention metadata, and expert routing metadata for the step. -4. **Forward pass.** `ContinuousBatchingEngine.step()` builds the batch, checks `_can_delegate_speculative(batch)`, and either enters `_step_speculative()` or falls back to `_execute_batch()` + sampling. The normal path still reaches `DistributedExpertExecutor.dispatch_local()` through the model blocks, so the same expert dispatch and prefetch hooks stay in play. -5. **Speculative delegation.** If the batch is a fresh singleton greedy prefill request, `_step_speculative()` delegates to `DFlashSpeculator.generate()` and then records committed counts through `Scheduler.update_after_step(..., committed_counts=...)`. -6. **Rollback bookkeeping.** `SpecDecodeState` and `apply_verify_step()` are low-level tested helpers for committed-count and truncation math. The live serving engine uses `Scheduler.update_after_step(..., committed_counts=...)` with `PagedKVCache` directly; these helpers describe and verify the contract but are not the live integration point. +4. **Forward pass.** `ContinuousBatchingEngine.step()` starts eligible persistent sessions and executes incompatible rows through the standard sampler without changing their policy. +5. **Speculative rounds.** The driver drafts, registers two-dimensional verify demand, verifies admitted rows, and publishes unseen canonical commits one token at a time. +6. **Cache mode.** Stage 4a refreshes temporary dense state. Eligible default-off Stage 4b DeepSeek sessions append/truncate engine-owned MLA pages. `SpecDecodeState` checks logical accounting; `apply_verify_step()` remains test-facing executable specification. 7. **Streaming.** `StreamManager` pushes partial deltas to open SSE clients. The FastAPI response emits OpenAI-shaped chunks. 8. **Termination.** When a sequence finishes or the client disconnects, the scheduler releases its KV blocks, `abort_request()` clears callbacks and request state, and the cancel path updates request stats. @@ -225,7 +247,7 @@ Both paths share: | `moe_infinity.MoE`, `moe_infinity.OffloadEngine`, `moe_infinity.__version__` | Documented package surface | Users | The top-level class and package names are documented surfaces; individual methods have the lifecycle stated in their own rows or guides. | | `MoE.generate()` | Deprecated, pending removal | Existing synchronous callers | Emits `DeprecationWarning`. It remains documented for transition and current validation only, with no compatibility promise beyond transition documentation. | | `MoE.serve()`, `moe_infinity.entrypoints.openai.api_server_v2`, `moe_infinity.entrypoints.openai.protocol`, and the documented routes in `docs/serving.md` | Documented server surface | Operators and integrators | Request/response shapes and route behavior follow the documented server contract. | -| `moe_infinity.spec_decode.DFlashConfig`, `DFlashSpeculator`, `read_dflash_config`, `validate_pairing`, `glm_dflash_available`, `glm_dflash_drafter_for`, `validate_glm_pairing` | Experimental exported `spec_decode` surface | Power users and contributors | Exported for experimentation and integration work, but not a stable compatibility contract. Changes should still be reflected in docs/release notes. | +| `moe_infinity.spec_decode.DFlashConfig`, `DFlashSpeculator`, protocol/trace types, `SessionDriver`, `read_dflash_config`, and pairing helpers | Experimental exported `spec_decode` surface | Power users and contributors | Exported for experimentation and integration work. The documented trace field schema is a rollout evidence contract; execution class signatures are not a stable user API. | | `moe_infinity.spec_decode.dflash.validate_drafter`, `validate_drafter_module`, `bind_shared_weights` | Internal DFlash helpers | Contributors working in `dflash.py` | Module-level implementation details; exported only from `dflash.py`, not from the package root. | | `moe_infinity.engine.*`, `moe_infinity.serving.*`, `moe_infinity.memory.*`, `moe_infinity.models.*`, `moe_infinity.kernel.*`, `moe_infinity.distributed.*`, `moe_infinity.runtime.*` | Internal | Contributors | No compatibility promise. Import paths, class names, and helper signatures may change. | @@ -233,21 +255,19 @@ If a symbol is exported but not listed in the documented package surface row, tr ## 6. Future Work -- **Unify `engine/` and `serving/`.** Today the synchronous and async paths - are two independent schedulers with duplicate data structures (`Sequence` - vs `SequenceGroup`, `SchedulerOutput` vs `SchedulerOutput`, etc.). A future - refactor should make `MoE.generate()` a synchronous facade over the - continuous batching engine so there is only one scheduling code path. +- **Unify scheduler/lifecycle code.** DFlash semantics are shared, but the + deprecated sync and async serving paths still have separate schedulers and + duplicate request data structures. A later removal of `MoE.generate()` can + retire that remaining lifecycle duplication. - **Expand `distributed/` tests.** The distributed module has smoke tests only (see `tests/python/unit/test_distributed_smoke.py`). Deeper coverage requires a multi-process CUDA harness. - **Multi-node distributed inference.** The current distributed module only supports single-host multi-GPU via NCCL; cross-host RPC scaffolding exists but is not production-tested. -- **Block-diffusion (dflash) serving.** A design proposal for serving - block-diffusion LLMs on the offloading runtime — dflash kernels, PD-dflash - scheduling, and an expert-prefetch lookahead — lives in - [`docs/design/pd-dflash-moe-serving.md`](./docs/design/pd-dflash-moe-serving.md). +- **Widen Stage 4b only with evidence.** Sampled paged serving, preemptible MLA + pages, hybrid paged rollback, and real DeepSeek/Qwen pairs remain future work. + GPT-OSS also remains without an executor route-ahead path. ## 7. Where to Look When … diff --git a/CHANGELOG.md b/CHANGELOG.md index 9f171402..ab156bb2 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -7,7 +7,10 @@ All notable changes to MoE-Infinity will be documented in this file. ### Added - Documentation hub at `docs/README.md` for users, operators, contributors, and project-history readers. -- DFlash documentation that distinguishes direct batch-1 greedy and sampled draft/verify from the greedy-gated `MoE.generate` and serving integrations, explains the current batch>1 greedy-only constraint, and limits continuous-batching and route-ahead claims to validated paths. +- Unified DFlash `SessionDriver`/`SpecSession` protocol, capability-selected bare-HF and rich backends, shared trace evidence, and per-request sampling streams. +- Direct bare-HF batch-1/batch>1 greedy, sampled, and mixed-row execution with dense reconstruction, right-padded output, and `last_generated_lengths`. +- Stage 4a persistent serving sessions and default-off Stage 4b engine-owned DeepSeek V2/V3 MLA pages for eligible greedy batch-1 requests. +- No-download unified-execution benchmark/validator and compatibility assertions that fail closed on sampling, ordering, cache invariant, or ownership failures. ### Changed @@ -15,6 +18,7 @@ All notable changes to MoE-Infinity will be documented in this file. - GLM-5.2-FP8 keeps routed FP8 experts in the host store while non-routed FP8 weights dequantize to BF16 on load. - Root README is now a concise discovery surface and points readers to the docs hub, model compatibility, DFlash, serving, troubleshooting, architecture, and changelog. - Release notes are split out of README and tracked here instead of being presented as shipped releases. +- Package version is now derived from git tags by setuptools-scm and written to `moe_infinity/_version.py` at build time, replacing the manual `MOEINF_VERSION`, `NIGHTLY_BASE_VERSION`, and hardcoded `setup.py`/`__init__.py` version strings. ### Deprecated @@ -25,8 +29,12 @@ All notable changes to MoE-Infinity will be documented in this file. - Serving-path DFlash now truncates KV cache to the committed prefix after each verify step, so emitted and cached tokens stay aligned. - GPT-OSS resident-load path now materializes MXFP4 blocks, scales, router, biases, and attention sinks instead of leaving placeholder tensors in place. - GLM FP8 store and reload parity now stays stable across fresh stores and reloads. +- PyPI publishing for both stable (`publish.yml`) and nightly (`publish-test.yml`): stable releases now take their version from the pushed git tag instead of always publishing `0.0.1`, and nightly sdists carry their version in `PKG-INFO` so `pip install --pre moe-infinity` no longer fails with a `MetadataInconsistent` version mismatch on rebuild. ### Known Limitations -- Batch > 1 DFlash is greedy-only, requires a bare HuggingFace target, and sampled batch > 1 remains unsupported. +- `MoE.generate()` retains narrower compatibility behavior than the direct API and does not generally expose sampled batch > 1. +- Stage 4b is default-off, greedy batch-1 DeepSeek V2/V3 only, resident-only, and non-preemptible; all ineligible and hybrid/Qwen paths use Stage 4a. +- No real DeepSeek or Qwen DFlash pair is claimed. GPT-OSS has named valid pairs but no executor route-ahead path. +- Required GPU fixture skips/unavailability are not rollout success; these Unreleased notes do not claim a release or unrun GPU validation. - Multi-node distributed inference is still unsupported. diff --git a/MANIFEST.in b/MANIFEST.in index b95b7d72..0d6b5831 100644 --- a/MANIFEST.in +++ b/MANIFEST.in @@ -1,2 +1,15 @@ recursive-include core *.cpp *.h *.cc recursive-include extensions/kernel *.cu *.h + +# With setuptools-scm installed, the sdist file-finder defaults to including +# every git-tracked file. Prune large directories that are not needed to build +# or install the package from source, keeping the published sdist lean. +prune docs +prune benchmarks +prune examples +prune tests +prune docker +prune scripts +prune results +prune logs +prune .github diff --git a/README.md b/README.md index 5bec6d58..5cff89f3 100644 --- a/README.md +++ b/README.md @@ -40,8 +40,8 @@ This open-sourced version is HuggingFace-friendly and differs from the version r - **Fast.** Activation-aware expert caching, prefetching, tracing, fused CUDA kernels, CUDA graph capture, Marlin INT4 GEMM, and FP4/MXFP4 expert paths keep the hot path lean. - **HuggingFace-native.** The `MoE` class remains the current in-process synchronous API, but `MoE.generate()` emits `DeprecationWarning` and is scheduled for removal. Use `MoE.serve()` for continuous batching; it starts an async HTTP service and is not a drop-in in-process return API. - **Production serving.** OpenAI-compatible HTTP server with continuous batching, paged KV cache, request scheduling with preemption, streaming (SSE), runtime hot reload, watchdog/health monitoring, and crash-recovery logging. A prefix-cache flag and cache scaffolding exist, but the current OpenAI request path does not actively reuse cached prefixes; see [docs/serving.md](docs/serving.md#prefix-caching). -- **Acceleration-aware.** Automatically integrates with [FlashAttention](https://github.com/Dao-AILab/flash-attention) and [FlashInfer](https://flashinfer.ai/) when installed, with graceful fallback to built-in kernels. -- **DFlash.** The experimental direct speculator API supports batch-1 greedy and sampled draft/verify without a stable API promise; deprecated `MoE.generate()` and continuous serving delegate only greedy singleton requests. Batch>1 is currently greedy-only on the bare HuggingFace target path. Route-ahead is an executor-path capability, not evidence of a validated target/drafter pair; see the model-by-model status in [docs/dflash.md](docs/dflash.md#compatibility). +- **Acceleration-aware.** Automatically integrates with [FlashAttention](https://github.com/Dao-AILab/flash-attention) and uses FlashInfer where the selected standard paged-attention backend supports it, with graceful fallback to built-in kernels. DeepSeek MLA currently uses the correct PyTorch fallback and does not claim FlashInfer acceleration. +- **DFlash.** One session semantic core now covers direct, deprecated-sync, and serving draft/verify decisions. The experimental direct bare-HF API supports batch-1/batch>1 greedy, sampled, and mixed rows. Physical rich batching is capability-gated; unsupported MLA/hybrid wrappers run grouped per-request sessions. Serving keeps Stage 4a dynamic fallback, with default-off Stage 4b paged MLA limited to eligible greedy batch-1 DeepSeek V2/V3. Pairing and executor route-ahead evidence remain separate; see [docs/dflash.md](docs/dflash.md). - **Multi-GPU.** Single-server multi-GPU with round-robin expert distribution, per-GPU caching, and an in-memory N-way tensor-parallel shard loader; see [docs/multi-gpu.md](docs/multi-gpu.md) and [docs/troubleshooting.md](docs/troubleshooting.md). ## Supported Models @@ -123,7 +123,7 @@ sudo apt-get update && sudo apt-get install -y build-essential cmake ninja-build # 2. Build tools + PyTorch. Match PyTorch's CUDA build to your CUDA toolkit # (pick the index URL for your CUDA version from https://pytorch.org). -pip install "setuptools>=78.1.1,<82" wheel ninja py-cpuinfo +pip install "setuptools>=78.1.1,<82" "setuptools-scm>=8" wheel ninja py-cpuinfo pip install torch --index-url https://download.pytorch.org/whl/cu128 # 3. CUTLASS headers (header-only; no separate build required) @@ -157,7 +157,7 @@ Post-installation, MoE-Infinity will automatically use FlashAttention when avail ### Enable FlashInfer (Optional) -Install [FlashInfer](https://flashinfer.ai/) for optimized paged attention kernels during prefill and decode. +Install [FlashInfer](https://flashinfer.ai/) for optional optimized standard paged-attention kernels during prefill and decode. It does not currently accelerate DeepSeek MLA. ```bash # Install the FlashInfer Python package (JIT-compiles kernels to match your Torch/CUDA): @@ -168,7 +168,7 @@ pip install -e '.[flashinfer]' Check the [FlashInfer installation guide](https://docs.flashinfer.ai/installation.html) for prebuilt-wheel options matching specific CUDA/PyTorch versions. -Post-installation, MoE-Infinity will automatically detect and use FlashInfer when available. When FlashInfer is not installed, MoE-Infinity gracefully falls back to its built-in attention kernels with no behavior change. +Post-installation, MoE-Infinity will detect and use FlashInfer where the selected backend supports it. When FlashInfer is not installed, it falls back to built-in attention kernels with no behavior change. ## Usage and Examples @@ -263,7 +263,19 @@ model = MoE("zai-org/GLM-5.2-FP8", { ## DFlash -See [docs/dflash.md](docs/dflash.md) for the batch-1 greedy/sampled flow, the current batch>1 greedy-only constraint, and the compatibility matrix that separates DFlash pairing validation from route-ahead executor wiring. +See [docs/dflash.md](docs/dflash.md) for unified session semantics, direct +greedy/sampled/mixed batching, per-row RNG and scalar-generator correlation, +dense reconstruction, output padding and `last_generated_lengths`, grouped +versus physical rich execution, Stage 4a/4b ownership, and the separate pairing +versus executor evidence matrix. No real DeepSeek DFlash pair or GPT-OSS +executor route-ahead is implied. + +No-download rollout gate: + +```bash +python benchmarks/dflash/validate_unified_execution.py --fixture tiny \ + --require-cache-invariants --require-order-invariance +``` ### Benchmarking diff --git a/RELEASE.md b/RELEASE.md index a62676a4..d8eaa7e3 100644 --- a/RELEASE.md +++ b/RELEASE.md @@ -44,6 +44,25 @@ Stable releases are automated through GitHub Actions workflows in `.github/workf - `.github/workflows/publish-test.yml`: publishes nightly pre-release builds from `main` to PyPI. - `.github/workflows/build-test.yml`: build validation for pull requests. +### One-time PyPI setup (Trusted Publishing) + +Publishing authenticates via PyPI Trusted Publishing (OIDC); there are no PyPI +username/password/token secrets in the repository. The `pypa/gh-action-pypi-publish` +action reads credentials only from its `user`/`password` inputs (not from +`TWINE_*` env vars) and, with none supplied plus `id-token: write`, uses OIDC. + +Before the first successful publish, register a trusted publisher on PyPI once +per workflow, at `https://pypi.org/manage/project/moe-infinity/settings/publishing/`: + +- Owner: `EfficientMoE` +- Repository name: `MoE-Infinity` +- Workflow name: `publish-test.yml` (nightly) — then add a second, identical + entry with Workflow name `publish.yml` (stable) +- Environment: leave blank + +Until these are registered, the publish step fails with +`invalid-publisher: ... no corresponding publisher`. + ### Steps to Release a New Version To release a new version, such as version 1.0.0, follow this order: @@ -51,10 +70,18 @@ To release a new version, such as version 1.0.0, follow this order: - close the current `## [Unreleased]` section in `CHANGELOG.md` - move shipped notes into the dated release entry - recreate a fresh `## [Unreleased]` section at the top -1. Update Version: - - Update `moe_infinity/__init__.py` (`__version__ = "..."`) to the new stable version. - - Ensure `setup.py` remains `version=os.getenv("MOEINF_VERSION", "...")` and update the default fallback version there to match the new stable version. - - If needed, bump `NIGHTLY_BASE_VERSION` in `.github/workflows/publish-test.yml` to the next planned stable series so nightly dev builds sort correctly. +1. Versioning is automatic (setuptools-scm): + - Versions are derived from the git tag at build time (see `pyproject.toml` + `[tool.setuptools_scm]`) and written to `moe_infinity/_version.py`. There + are no version strings to edit in `setup.py` or `moe_infinity/__init__.py`. + - Tag `vX.Y.Z` publishes stable `X.Y.Z`. Every later commit on `main` + publishes as the next-patch pre-release `X.Y.(Z+1).devN`, so nightlies + always sort above the last stable and below the next one, with no manual + `NIGHTLY_BASE_VERSION` bump. + - First release only: the repository has no tags yet, so until the first tag + exists nightlies version as `0.1.devN`. Create the initial tag (for example + `git tag v0.0.1`) on the release commit to anchor the `0.0.x` series; + afterwards nightlies become `0.0.2.devN` automatically. 2. Review the release checklist above - confirm model and capability coverage - verify install and quick starts diff --git a/benchmarks/dflash/unified_execution_benchmark.py b/benchmarks/dflash/unified_execution_benchmark.py new file mode 100644 index 00000000..3b762097 --- /dev/null +++ b/benchmarks/dflash/unified_execution_benchmark.py @@ -0,0 +1,260 @@ +#!/usr/bin/env python3 +"""No-download DFlash unified-execution benchmark and evidence reporter. + +The ``tiny`` fixture measures CPU protocol operations. It does not estimate +real-checkpoint throughput and leaves unsupported production capabilities +explicit rather than inventing values. +""" + +from __future__ import annotations + +import argparse +import json +import math +import os +import sys +import time +from collections import Counter +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +if os.environ.get("MKL_THREADING_LAYER") == "INTEL": + os.environ["MKL_THREADING_LAYER"] = "GNU" + +import torch + +ROOT = Path(__file__).resolve().parents[2] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from moe_infinity.serving.mla_cache import MLAPagedKVCache +from moe_infinity.serving.spec_cache_adapter import PagedCacheAdapter +from moe_infinity.spec_decode._dflash_sample_ops import acceptance_sampled +from moe_infinity.spec_decode.protocols import ExecutorEvidence, PairingEvidence + + +def _milliseconds(start: float) -> float: + return (time.perf_counter() - start) * 1_000.0 + + +def _distribution(counts: Counter[int], samples: int, size: int) -> list[float]: + return [counts[index] / samples for index in range(size)] + + +def _tvd(left: list[float], right: list[float]) -> float: + return 0.5 * sum(abs(a - b) for a, b in zip(left, right)) + + +def _kl(left: list[float], right: list[float]) -> float: + epsilon = 1e-12 + return sum( + a * math.log((a + epsilon) / (b + epsilon)) + for a, b in zip(left, right) + if a > 0 + ) + + +@dataclass(frozen=True) +class SampledLawMeasurements: + tvd_value: float + kl_value: float + sample_count: int + round_count: int + accepted_drafts: int + committed_tokens: int + rollback_count: int + elapsed_seconds: float + + +def _sampled_law(samples: int = 2_000) -> SampledLawMeasurements: + proposal = torch.tensor([[0.70, 0.30]], dtype=torch.float64) + target = torch.tensor([[0.25, 0.75], [0.60, 0.40]], dtype=torch.float64) + speculative: Counter[int] = Counter() + reference: Counter[int] = Counter() + accepted = 0 + committed = 0 + rollback_count = 0 + round_count = 0 + start = time.perf_counter() + for seed in range(samples): + spec_generator = torch.Generator().manual_seed(seed) + draft = torch.multinomial( + proposal[0], 1, generator=spec_generator + ).reshape(1) + decision = acceptance_sampled( + proposal, target, draft, generator=spec_generator + ) + token = int(draft[0]) if decision.accept else int(decision.final_token) + speculative[token] += 1 + accepted += int(decision.accept) + committed += 1 + int(decision.accept) + rollback_count += int(not decision.accept) + round_count += 1 + + reference_generator = torch.Generator().manual_seed(100_000 + seed) + reference[ + int(torch.multinomial(target[0], 1, generator=reference_generator)) + ] += 1 + elapsed = time.perf_counter() - start + spec_dist = _distribution(speculative, samples, 2) + ref_dist = _distribution(reference, samples, 2) + return SampledLawMeasurements( + tvd_value=_tvd(spec_dist, ref_dist), + kl_value=_kl(spec_dist, ref_dist), + sample_count=sum(speculative.values()), + round_count=round_count, + accepted_drafts=accepted, + committed_tokens=committed, + rollback_count=rollback_count, + elapsed_seconds=elapsed, + ) + + +def _order_invariance() -> bool: + seeds = {"a": 17, "b": 29, "c": 41} + + def run(order: tuple[str, ...]) -> dict[str, tuple[int, ...]]: + generators = { + name: torch.Generator().manual_seed(seed) + for name, seed in seeds.items() + } + rows: dict[str, list[int]] = {name: [] for name in seeds} + probabilities = torch.tensor([0.2, 0.3, 0.5]) + for _ in range(16): + for name in order: + rows[name].append( + int( + torch.multinomial( + probabilities, 1, generator=generators[name] + ) + ) + ) + return {name: tuple(tokens) for name, tokens in rows.items()} + + return run(("a", "b", "c")) == run(("c", "a", "b")) + + +def _cache_evidence() -> dict[str, Any]: + cache = MLAPagedKVCache( + num_blocks=16, + block_size=2, + num_layers=1, + latent_dim=2, + rope_dim=2, + dtype=torch.float32, + device=torch.device("cpu"), + ) + first = PagedCacheAdapter(cache, seq_id=1, initial_length=3) + second = PagedCacheAdapter(cache, seq_id=2, initial_length=2) + second_pages = tuple(cache.get_block_table(2)) + snapshot = first.snapshot() + first.append(4) + peak = len(cache.get_block_table(1)) + len(cache.get_block_table(2)) + first.truncate(5) + first.restore(snapshot) + isolated = tuple(cache.get_block_table(2)) == second_pages + invariant = first.logical_length() == 3 and second.logical_length() == 2 + + cancel_start = time.perf_counter() + first.release() + cancellation_latency = _milliseconds(cancel_start) + try: + cache.get_block_table(1) + released = False + except KeyError: + released = True + second.release() + return { + "cache_pages_peak": peak, + "cache_invariants": invariant, + "ownership_isolation": isolated, + "cancellation_released_pages": released, + "cancellation_latency_ms": cancellation_latency, + "preemption_policy": "resident-only no-preempt for paged MLA", + "preemption_recovery": "not exercised: swap_out/swap_in intentionally return false", + } + + +def run_tiny() -> dict[str, Any]: + prompt = torch.arange(32, dtype=torch.float32) + weights = torch.arange(32 * 16, dtype=torch.float32).reshape(32, 16) + prefill_start = time.perf_counter() + for _ in range(64): + _ = prompt @ weights + prefill_ms = _milliseconds(prefill_start) / 64.0 + + sampled = _sampled_law() + verify_start = time.perf_counter() + for seed in range(128): + generator = torch.Generator().manual_seed(seed) + _ = acceptance_sampled( + torch.tensor([[0.7, 0.3]]), + torch.tensor([[0.25, 0.75], [0.6, 0.4]]), + torch.tensor([seed % 2]), + generator=generator, + ) + verify_ms = _milliseconds(verify_start) / 128.0 + + cache = _cache_evidence() + pairing = PairingEvidence( + failure_reason="tiny fixture has no checkpoint pair" + ) + executor = ExecutorEvidence( + fallback_reason="tiny fixture has no expert executor" + ) + report: dict[str, Any] = { + "fixture": "tiny", + "measurement_scope": "synthetic no-checkpoint CPU fixture", + "prefill_latency_ms": prefill_ms, + "verify_latency_ms": verify_ms, + "decode_elapsed_seconds": sampled.elapsed_seconds, + "decode_committed_tokens_per_second": ( + sampled.committed_tokens / sampled.elapsed_seconds + ), + "sample_count": sampled.sample_count, + "round_count": sampled.round_count, + "accepted_drafts": sampled.accepted_drafts, + "committed_tokens": sampled.committed_tokens, + "rollback_count": sampled.rollback_count, + "replay_count": 0, + "rng_order_invariant": _order_invariance(), + "sampled_tvd_value": sampled.tvd_value, + "sampled_kl_value": sampled.kl_value, + "metric_units": { + "prefill_latency_ms": "milliseconds per prefill operation", + "verify_latency_ms": "milliseconds per verify operation", + "decode_elapsed_seconds": "seconds", + "decode_committed_tokens_per_second": "committed tokens per second", + "sampled_tvd_value": "dimensionless", + "sampled_kl_value": "nats", + "cache_pages_peak": "pages", + "cancellation_latency_ms": "milliseconds", + }, + "execution_mode": "tiny_cpu_protocol_fixture", + "pairing_evidence": pairing.as_dict(), + "executor_evidence": executor.as_dict(), + "route_attempted_layers": [], + "route_fired_layers": [], + "per_request_rich_calls": 0, + "physical_rich_calls": 0, + **cache, + } + return report + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--fixture", choices=("tiny",), default="tiny") + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + report = run_tiny() + if args.json: + print(json.dumps(report, sort_keys=True)) + else: + print(json.dumps(report, indent=2, sort_keys=True)) + return 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/benchmarks/dflash/validate_unified_execution.py b/benchmarks/dflash/validate_unified_execution.py new file mode 100644 index 00000000..800c9dd3 --- /dev/null +++ b/benchmarks/dflash/validate_unified_execution.py @@ -0,0 +1,128 @@ +#!/usr/bin/env python3 +"""Fail-closed local rollout gates for unified DFlash execution.""" + +from __future__ import annotations + +import argparse +import json +import os +import sys +from pathlib import Path +from typing import Any + +if os.environ.get("MKL_THREADING_LAYER") == "INTEL": + os.environ["MKL_THREADING_LAYER"] = "GNU" + +import torch + +ROOT = Path(__file__).resolve().parents[2] +if str(ROOT) not in sys.path: + sys.path.insert(0, str(ROOT)) + +from benchmarks.dflash.unified_execution_benchmark import run_tiny + + +def validate( + *, + require_cache_invariants: bool, + require_order_invariance: bool, + require_gpu: bool, +) -> tuple[dict[str, Any], bool]: + benchmark = run_tiny() + gpu_available = torch.cuda.is_available() + gpu_fixture_enabled = os.environ.get("MOE_DFLASH_GPU") == "1" + gpu_readiness_pass = gpu_available and gpu_fixture_enabled + sampled_tvd_value = float(benchmark["sampled_tvd_value"]) + sampled_kl_value = float(benchmark["sampled_kl_value"]) + gates = { + "cache_invariants": bool(benchmark["cache_invariants"]), + "ownership_isolation": bool(benchmark["ownership_isolation"]), + "order_invariance": bool(benchmark["rng_order_invariant"]), + "sampled_tvd_value": sampled_tvd_value, + "sampled_tvd_pass": sampled_tvd_value <= 0.10, + "sampled_kl_value": sampled_kl_value, + "sampled_kl_pass": sampled_kl_value <= 0.05, + "pairing_executor_separate": set( + benchmark["pairing_evidence"] + ).isdisjoint({"wiring_reachable", "attempted_layers", "fired_layers"}), + "paged_ownership_released": bool( + benchmark["cancellation_released_pages"] + ), + } + required = [ + gates["sampled_tvd_pass"], + gates["sampled_kl_pass"], + gates["pairing_executor_separate"], + gates["paged_ownership_released"], + gates["ownership_isolation"], + ] + if require_cache_invariants: + required.append(gates["cache_invariants"]) + if require_order_invariance: + required.append(gates["order_invariance"]) + if require_gpu: + required.append(gpu_readiness_pass) + + passed = all(required) + report: dict[str, Any] = { + "status": "PASS" if passed else "FAIL", + "fixture": "tiny", + "checkpoint_downloads": False, + "required_gpu_fixture": require_gpu, + "gpu_readiness_required": require_gpu, + "gpu_readiness_pass": gpu_readiness_pass, + "gpu_gate_kind": "readiness only", + "gpu_harness_executed": False, + "gpu_harness_command": ( + "MOE_DFLASH_GPU=1 CUDA_VISIBLE_DEVICES=0 pytest -q " + "tests/python/dflash/test_gpu_20b_dflash.py " + "tests/python/dflash/test_gpu_serving_dflash.py -m gpu" + ), + "gpu_available": gpu_available, + "gpu_fixture_enabled": gpu_fixture_enabled, + **gates, + "compatibility": { + "pairing_evidence": benchmark["pairing_evidence"], + "executor_evidence": benchmark["executor_evidence"], + "execution_mode": benchmark["execution_mode"], + }, + "trace_summary": { + key: benchmark[key] + for key in ( + "accepted_drafts", + "committed_tokens", + "sample_count", + "round_count", + "rollback_count", + "replay_count", + "per_request_rich_calls", + "physical_rich_calls", + ) + }, + } + return report, passed + + +def main() -> int: + parser = argparse.ArgumentParser() + parser.add_argument("--fixture", choices=("tiny",), default="tiny") + parser.add_argument("--require-cache-invariants", action="store_true") + parser.add_argument("--require-order-invariance", action="store_true") + parser.add_argument("--require-gpu", action="store_true") + parser.add_argument("--json", action="store_true") + args = parser.parse_args() + report, passed = validate( + require_cache_invariants=args.require_cache_invariants, + require_order_invariance=args.require_order_invariance, + require_gpu=args.require_gpu, + ) + print( + json.dumps(report, sort_keys=True) + if args.json + else json.dumps(report, indent=2, sort_keys=True) + ) + return 0 if passed else 1 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/core/parallel/expert_dispatcher.cpp b/core/parallel/expert_dispatcher.cpp index f3b525a1..ea6ddc8e 100644 --- a/core/parallel/expert_dispatcher.cpp +++ b/core/parallel/expert_dispatcher.cpp @@ -131,6 +131,9 @@ ExpertDispatcher::ExpertDispatcher(int num_experts, int num_layers, int dtype, cudaStream_t exec_stream; cudaStreamCreateWithFlags(&exec_stream, cudaStreamNonBlocking); exec_streams_.emplace_back(exec_stream); + cudaEvent_t exec_done_event; + cudaEventCreateWithFlags(&exec_done_event, cudaEventDisableTiming); + exec_done_events_.emplace_back(exec_done_event); modules_[i] = new MoEMLP(dtype, expert_type); @@ -616,6 +619,17 @@ void ExpertDispatcher::GPUExecFunc(int gpu_id, int thread_idx) { auto device = CUDA_DEVICE(gpu_id); auto expert_idx = args.expert_node->expert_idx; + // Enter the exec stream before any tensor op and order it after the + // producer stream that wrote hidden_states_/router_mask_ (recorded in + // SetInputs). exec streams are non-blocking, so without this the + // gather/copy below could run against not-yet-written input memory. + c10::cuda::CUDAStream torch_stream = + c10::cuda::getStreamFromExternal(stream, gpu_id); + c10::cuda::CUDAStreamGuard guard(torch_stream); + if (input_ready_event_ != nullptr) { + cudaStreamWaitEvent(stream, input_ready_event_, 0); + } + auto token_mask = router_mask_.index({"...", expert_idx}); torch::Tensor input = (batch_size == 1) ? hidden_states_.to(device) @@ -635,10 +649,6 @@ void ExpertDispatcher::GPUExecFunc(int gpu_id, int thread_idx) { } } - c10::cuda::CUDAStream torch_stream = - c10::cuda::getStreamFromExternal(stream, gpu_id); - c10::cuda::CUDAStreamGuard guard(torch_stream); - if (expert_type_ == GPT_OSS_MOE_DENSE_ACT_DENSE) { modules_[thread_idx]->DequantMxfp4Params(stream); } @@ -721,6 +731,7 @@ std::vector ExpertDispatcher::Wait() { std::unique_lock lock(pending_mutex_); pending_cv_.wait(lock, [&] { return pending_.load() == 0; }); + SyncExecStreamsWithCurrent(); num_enqueued_.store(0); std::vector output_queue; @@ -738,6 +749,10 @@ torch::Tensor ExpertDispatcher::WaitHiddenStates() { #endif std::unique_lock lock(pending_mutex_); pending_cv_.wait(lock, [&] { return pending_.load() == 0; }); + // pending_ hits zero when the accumulation is enqueued, not complete; + // order the caller's stream after every exec stream before the result + // tensor is consumed. + SyncExecStreamsWithCurrent(); num_enqueued_.store(0); return final_hidden_states_; } @@ -752,4 +767,17 @@ void ExpertDispatcher::SetInputs(const torch::Tensor& hidden_states, router_mask_ = router_mask; router_weight_ = router_weight; // this can be float32 final_hidden_states_ = torch::zeros_like(hidden_states, options); + + if (input_ready_event_ == nullptr) { + cudaEventCreateWithFlags(&input_ready_event_, cudaEventDisableTiming); + } + cudaEventRecord(input_ready_event_, c10::cuda::getCurrentCUDAStream()); +} + +void ExpertDispatcher::SyncExecStreamsWithCurrent() { + auto current = c10::cuda::getCurrentCUDAStream(); + for (size_t i = 0; i < exec_streams_.size(); ++i) { + cudaEventRecord(exec_done_events_[i], exec_streams_[i]); + cudaStreamWaitEvent(current.stream(), exec_done_events_[i], 0); + } } diff --git a/core/parallel/expert_dispatcher.h b/core/parallel/expert_dispatcher.h index 47d3b69a..8021f7ad 100644 --- a/core/parallel/expert_dispatcher.h +++ b/core/parallel/expert_dispatcher.h @@ -91,6 +91,12 @@ class ExpertDispatcher : public base::noncopyable { for (auto& stream : fetch_streams_) { cudaStreamDestroy(stream); } + for (auto& event : exec_done_events_) { + cudaEventDestroy(event); + } + if (input_ready_event_ != nullptr) { + cudaEventDestroy(input_ready_event_); + } for (auto* m : modules_) { delete m; } @@ -171,7 +177,11 @@ class ExpertDispatcher : public base::noncopyable { std::mutex output_mutex_; std::mutex accum_mutex_; + void SyncExecStreamsWithCurrent(); + std::vector exec_streams_; + std::vector exec_done_events_; + cudaEvent_t input_ready_event_ = nullptr; std::vector fetch_streams_; std::unique_ptr[]> gpu_overload_; diff --git a/docs/dflash.md b/docs/dflash.md index 5af1b342..bba8fc78 100644 --- a/docs/dflash.md +++ b/docs/dflash.md @@ -1,256 +1,218 @@ -# DFlash speculative decoding +# DFlash unified execution -MoE-Infinity ships a native DFlash draft, verify, rollback path for GPT-OSS and other supported MoE models. The deprecated synchronous wrapper `MoE.generate(..., speculative_draft=...)` remains the current in-process integration and emits `DeprecationWarning`; `MoE.serve(..., speculative_draft=...)` is the recommended continuous-batching HTTP path and is not a drop-in return API. Direct `DFlashSpeculator.generate(...)` is an experimental alternative for custom harnesses, with no stable API promise, and supports batch-1 greedy and sampled decoding. The `MoE` integrations are greedy-gated and delegate only singleton requests in the server. Batch > 1 DFlash stays on the bare HuggingFace target path. +> **Security warning:** DFlash drafter checkpoints and tokenizer/model code may +> require `trust_remote_code=True`. That setting permits arbitrary Python code +> from the remote repository to run during loading. Use it only with a trusted, +> pinned drafter revision, and review the revision before loading it. See the +> model compatibility matrix for the pairing evidence boundary. -For server-side gating, sampled-request fallback, and exact troubleshooting language, see [Serving](serving.md), [Architecture](../ARCHITECTURE.md), and [Troubleshooting](troubleshooting.md). +DFlash now has **one semantic core** for draft, verify, acceptance, rollback, +sampling, stop handling, and traces. Direct generation and the deprecated sync +facade use `SessionDriver` with request-scoped sessions; serving uses the same +canonical `SpecSession` transitions through its lifecycle-owning +`SpecSessionDriver`. `_generate_batched` is only an argument/output adapter. -## Quick start +This is an implementation statement, not a claim that every model, cache, or +serving mode has the same capabilities. Backend declarations gate physical +batching, sampling, rich forwarding, route-ahead, and cache ownership. -This example requires a CUDA-capable source installation, access to both -checkpoints, and enough GPU cache, host memory, and SSD capacity for GPT-OSS-120B. -Checkpoint loading may download weights when they are not already cached. For a -configurable runnable script using the same defaults, see -[`examples/dflash_gpt_oss_example.py`](../examples/dflash_gpt_oss_example.py). -The drafter loads with `trust_remote_code=True`, which executes code from the -checkpoint repository. Use only a trusted drafter and pin the repository -revision for reproducible or security-sensitive deployments. +## User surfaces -```python -from transformers import AutoTokenizer - -from moe_infinity import MoE -from moe_infinity.spec_decode import DFlashSpeculator - -target = "openai/gpt-oss-120b" -drafter = "z-lab/gpt-oss-120b-DFlash" -prompt = "Question: What is the capital of France?\nAnswer:" +| Surface | Delivered behavior | Important limits | +| --- | --- | --- | +| `DFlashSpeculator.generate` | Bare-HF batch 1 and batch > 1 support greedy, sampled, and mixed greedy and sampled rows. One session is created per row. | Experimental API. A target/drafter pair is still required. | +| Direct rich target, batch 1 | Canonical per-request rich execution. | Model-specific caches still apply. | +| Direct rich target, batch > 1 | A wrapper declaring the complete row-aware capability can use a physically batched rich forward. Otherwise rows use grouped per-request sessions. | Grouped per-request execution is not physically batched execution. DeepSeek MLA and hybrid/Qwen wrappers currently fall back per request. | +| `MoE.generate(..., speculative_draft=...)` | Compatibility facade over the same session semantics. Greedy batch 1 and eligible greedy batch > 1 retain tensor compatibility. | MoE.generate() is deprecated and emits `DeprecationWarning`. Sampled batch > 1 is not widened by this facade; Qwen3.5 rejects non-greedy speculative use explicitly. | +| Continuous serving | Persistent per-sequence sessions, streaming commits, cancellation cleanup, and scheduler-controlled verify admission. | Stage 4a and Stage 4b cache modes have different eligibility; see below. Do not infer sampled paged serving from direct sampled support. | -tokenizer = AutoTokenizer.from_pretrained(target) -input_ids = tokenizer(prompt, return_tensors="pt").input_ids +## Direct API -model = MoE(target, { - "offload_path": "/ssd/moe-infinity/gpt-oss-120b", - "device_memory_ratio": 0.75, -}) -spec = DFlashSpeculator(model, drafter) +```python +import torch +from moe_infinity.spec_decode import DFlashSpeculator -output_ids = model.generate( +# target and drafter are already constructed; no checkpoint loading is shown. +spec = DFlashSpeculator.from_models(target, drafter, config=config, device="cpu") +output = spec.generate( input_ids, - max_new_tokens=64, - do_sample=False, - speculative_draft=spec, + attention_mask=attention_mask, + max_new_tokens=[16, 24], + temperature=[0.0, 0.8], + top_p=[1.0, 0.9], + generator=[None, torch.Generator().manual_seed(7)], ) ``` -If you want sampled batch-1 DFlash, call `spec.generate(...)` directly. The -`MoE.generate` and serving integrations stay on the standard path for sampled -requests today. The Qwen3.5-MoE wrapper is the explicit exception: sampled -`MoE.generate(..., speculative_draft=...)` raises `ValueError` there, because -that model path requires greedy speculative decode. +### Batched inputs and outputs + +- Bare-HF batches may combine greedy and sampled rows. Budgets, stop sets, + temperatures, top-k, top-p, and generators are row-local. +- Ragged prompts must be left-padded with a 0/1 `attention_mask`; each row must + end in a real token. Equal-length prompts may omit the mask. +- Results form a right-padded tensor. `last_generated_lengths` records the true + generated length of every row so callers can ignore output padding. +- Direct dense caches use lockstep physical lengths where Hugging Face requires + a rectangle. A row ahead of the shared cache may re-feed already selected + tokens for **dense cache reconstruction**. Re-fed tokens are neither emitted + twice nor sampled again, and retained proposal distributions stay intact. + +### Request RNG contract + +Every sampled session owns a per-row generator and retains the drafter proposal +distribution for every slot until verification. Row order or the presence of an +unrelated row therefore does not change an explicitly row-seeded request. + +Passing one scalar generator to batch > 1 clones the same initial state for +each row. Identical requests can consequently be **correlated**. Pass one +generator per row for independent explicit streams. Results are not bit-exact +across batch shapes: the guarantee is row order/composition invariance for a +fixed request and row-local stream, not equality between all physical batch +layouts. In short, outputs are not bit-exact across batch shapes. + +## Rich execution capability + +`RichBatchMetadata` carries row offsets and lengths, masks, positions, cache +handles, request contexts, and route contexts. `RichForwardResult` carries +row-aligned logits, hidden states, and cache handles. Physical rich batching is +enabled only when the wrapper's row-aware capability guard declares the full +contract. Trace/benchmark evidence distinguishes `per_request_rich_calls` from +`physical_rich_calls`. + +Models with MLA or hybrid cache layouts currently keep the conservative +grouped per-request fallback unless their exact cache contract is supported. +This includes the current Qwen/hybrid fallback. Qwen evidence is tiny-fixture +only; no real Qwen target/drafter validation is claimed. + +## Serving cache modes + +### Stage 4a: `temporary_dynamic` + +Stage 4a is the compatibility mode. Each session temporarily owns a private +target/draft temporary DynamicCache while `ContinuousBatchingEngine` continues to +own scheduling and lifecycle. Ineligible Stage 4b requests remain here. A +sampled serving request using this fallback is not evidence of sampled serving +through paged MLA, and no sampled paged-serving claim is made. + +### Stage 4b: `paged_mla` + +Task 8.5 restored the DeepSeek MLA prerequisite before Task 9. Stage 4b is a +default-off target selected by `enable_deepseek_mla_paging=True`, and only for +eligible greedy batch-1 DeepSeek V2/V3 MLA sessions. Each eligible request has +exactly one engine-owned target paged store: standard `PagedKVCache` for +standard attention, or packed `MLAPagedKVCache` for DeepSeek MLA. +`PagedCacheAdapter` supplies per-sequence append, snapshot, truncate, attention +metadata, and release. The **draft cache remains separate** and never owns the +target allocation. + +The resident admission guard is implemented before `paged_mla` selection. It +caps active paged sessions with +`max_resident_paged_speculative_sessions=1` by default and requires prompt +plus declared output capacity, plus up to `DFlash block_size - 1` transient +verify tokens, to leave +`min_free_mla_blocks_after_admission=1` free MLA block by default. Existing +paged sessions' unallocated committed and transient headroom is included. All +demand is block-rounded using the MLA page size. Cap or reserve rejection records a structured reason/counter and +immediately starts the same request in Stage 4a `temporary_dynamic`; there is +no scheduler wait loop and no sampling downgrade. Releasing an admitted +session removes it from the active count, allowing a later request to qualify. +The dense Stage 4a target cache can increase total GPU memory use despite not +using MLA pages. Failed adapter/session construction is counted as +`begin_failed`, never `admitted`. + +DeepSeek MLA currently has a resident-only, non-preemptible policy: +`MLAPagedKVCache` has no swap/resume implementation. All DRAFT/VERIFY +speculative sessions, including Stage 4a sessions using a temporary +`DynamicCache`, remain resident and are not preempted while in flight. This +uses extra GPU memory and can increase wait time for ordinary requests. The +implemented cap/reserve guard bounds new paged admission and avoids starvation +by silent waiting; it is not a general fairness proof. It reserves the declared +request budget but cannot prevent unrelated cache consumers from exhausting +pages; the affected request is cleaned up and the current engine step re-raises +that allocator failure. Stage 4a also +temporarily double-allocates target state while its private cache is live. There +is no claim that speculative sessions can be swapped out and resumed, and +preemption is not implemented for these in-flight sessions. +`swap_out()` and +`swap_in()` intentionally return false, and the scheduler does not preempt +in-flight DRAFT/VERIFY sessions. This is not hybrid paged rollback, and Qwen or +other hybrid models remain Stage 4a fallbacks. Cancellation and completion +release the sequence's pages without touching another sequence. + +DeepSeek MLA does not currently use FlashInfer acceleration; it uses the +correct PyTorch fallback. Installing FlashInfer does not change that scope. + +There is no real DeepSeek DFlash pair validated in this repository. Stage 4b +tests prove cache ownership and DeepSeek MLA adapter behavior with tiny/local +models; they do not establish a production target/drafter checkpoint pair. + +## Pairing and executor evidence + +Pairing and execution are independent dimensions: + +- `pairing_evidence` describes config, dimensions, vocabulary, mask token, + target layers, block constraints, module checks, and any named checkpoint + scope. +- `executor_evidence` describes executor reachability, attempted/fired + route-ahead layers, actual expert unions, bytes, coverage, and fallback. + +GPT-OSS-20B and GPT-OSS-120B have valid published target/drafter pairs in the +repo's evidence, but the resident GPT-OSS expert path has no executor +route-ahead route. Conversely, executor wiring on DeepSeek/Qwen/Mixtral does not +create a valid DFlash pair. Route-ahead remains observer-only and cannot change +tokens, acceptance, or cache state. + +## Trace contract + +The direct and serving paths share the logical `SessionTrace` schema: +`request_id`, `backend`, `cache_kind`, sampled mode, `round_count`, `accepted`, +`committed`, `emitted`, `rollback`, `replay`, finish reason, route status, +`pairing_evidence`, and `executor_evidence`. These fields are the stable rollout +evidence schema; low-level cache objects remain internal. + +## Validation and benchmarks + +No-download CPU gates: + +```bash +pytest -q tests/python/dflash/test_compatibility_matrix.py +python benchmarks/dflash/unified_execution_benchmark.py --fixture tiny +python benchmarks/dflash/validate_unified_execution.py --fixture tiny \ + --require-cache-invariants --require-order-invariance +pytest -q tests/python/dflash +pytest -q tests/python/serving +``` + +The tiny benchmark labels itself as a synthetic no-checkpoint CPU fixture. It +reports measured prefill/verify/decode work, actual sample/round/committed-token +counts, observed rollback/replay events, RNG/order invariance, +`sampled_tvd_value` (dimensionless), `sampled_kl_value` (nats), cache pages, +execution mode, separate pairing/executor evidence, and rich-call counts. It +does not project those timings onto a checkpoint. + +GPU gates are opt-in and may download nothing unless checkpoints are already +cached: + +```bash +MOE_DFLASH_GPU=1 CUDA_VISIBLE_DEVICES=0 \ + pytest -q tests/python/dflash/test_gpu_20b_dflash.py \ + tests/python/dflash/test_gpu_serving_dflash.py -m gpu +MOE_DFLASH_GPU=1 CUDA_VISIBLE_DEVICES=0 \ + pytest -q tests/python/dflash/test_gpu_120b.py -m gpu +python benchmarks/dflash/validate_unified_execution.py --fixture tiny \ + --require-cache-invariants --require-order-invariance --require-gpu +``` -## Capability matrix +`--require-gpu is a readiness gate`: it checks only that CUDA is available and +the fixture environment is enabled. It does not execute the GPU harness. The +actual GPU pytest command remains separate and required. A skipped or +unavailable required fixture is not success; do not mark the GPU rollout gate +complete unless that command actually ran and passed. -| Capability | Status | Evidence or constraints | -| --- | --- | --- | -| Batch-1 direct `spec.generate`, greedy | Implemented, validated | `tests/python/dflash/test_native_step.py::test_native_multistep_greedy_matches_plain_greedy`, `tests/python/dflash/test_edge_cases.py` | -| Batch-1 direct `spec.generate`, sampled | Implemented, validated on CPU tiny fixtures | `tests/python/dflash/test_sampled_spec.py::test_sampled_generate_is_seed_deterministic`; this path is direct, not through `MoE.generate` | -| Batch > 1 direct `spec.generate`, greedy | Implemented, validated | `tests/python/dflash/test_batched_spec.py::test_batched_matches_looped_singles_token_identical` | -| `MoE.generate(..., speculative_draft=...)`, greedy batch-1 | Implemented, validated | `tests/python/dflash/test_engine_wire.py`, `tests/python/dflash/test_spec_seam.py` | -| `MoE.generate(..., speculative_draft=...)`, sampled batch-1 | Standard path for most models; Qwen3.5-MoE raises `ValueError` | The native speculator is greedy-gated in the engine, and the Qwen3.5 wrapper rejects sampled speculative decode on this path | -| Continuous-batching serving | Implemented, validated on GPT-OSS-20B | `tests/python/dflash/test_gpu_serving_dflash.py`; only a fresh singleton greedy request can delegate | -| Route-ahead expert prefetch | Implemented, scheduling-only, validated on synthetic offloaded shells | `tests/python/dflash/test_route_ahead_metrics.py`, `tests/python/dflash/test_route_ahead_wire.py`, `tests/python/dflash/test_qwen35_hybrid_rollback.py`; gpt-oss does not wire the executor | - -## Configuration - -### Direct speculator API - -- `DFlashSpeculator(moe, draft_model_path, device=None, dtype=torch.bfloat16)` - loads the drafter with `trust_remote_code=True` and defaults to bfloat16. -- `DFlashSpeculator.from_models(moe, draft_model, config=None, device=None)` - skips checkpoint loading and reuses an already-built drafter module. -- `DFlashSpeculator.generate(input_ids, max_new_tokens=256, temperature=0.0, - stop_token_ids=None, top_k=0, top_p=1.0, attention_mask=None)` is the public - experimental decode entry point; it has no stable API guarantee. -- `DFlashSpeculator.enable_route_ahead_stats()` creates or resets a read-only - recorder. `route_ahead_stats` starts as `None`. -- `read_dflash_config(draft_hf_config)` requires `block_size`, `mask_token_id`, - `target_layer_ids`, `hidden_size`, and `vocab_size`. `num_target_layers` is - optional and defaults to `-1`. -- `validate_pairing(...)` checks `hidden_size`, `vocab_size`, `mask_token_id < - vocab_size`, `block_size >= 2`, and target layer range. -- `validate_drafter(...)`, `validate_drafter_module(...)`, and - `bind_shared_weights(...)` are internal helpers inside `dflash.py`; they are - used by the implementation and tests, but they are not exported from - `moe_infinity.spec_decode.__init__`. - -### MoE and server knobs - -- `MoE.generate(..., speculative_draft=...)` attaches the speculator only for - the current call. Omit the kwarg, or pass `None` or `False`, to detach it. - For Qwen3.5-MoE, sampled requests on this path raise `ValueError` instead of - falling back silently. -- Greedy DFlash delegation is gated by `temperature=0`, `top_k=0`, and - `top_p=1.0`; `do_sample=False` is the usual caller intent for that path, but - it is not a separate runtime gate. -- `MoE.serve(..., speculative_draft=None)` defaults to `device_memory_ratio= - 0.75`, `kv_cache_ratio=0.25`, `max_batch_size=32`, and - `enable_prefix_caching=False`. -- `python -m moe_infinity.entrypoints.openai.api_server_v2 --speculative-draft - ` exposes the same path from the CLI. The parser default is off. - -If a sampled serving request falls back to the standard path, that is expected behavior, not a pairing failure; the pairing check is a separate validation step. - -## Draft, Verify, Commit, and Rollback - -1. **Draft** - - Build a block of the form `[anchor, MASK, MASK, ...]`. - - The target hidden states at the configured `target_layer_ids` are - concatenated into the drafter context feature. - - The drafter and target share the target `embed_tokens` and `lm_head`. - -2. **Verify** - - The target runs one full-logits forward over the whole block. - - Greedy mode uses the target argmax agreement rule. - - Sampled mode uses warped draft and target probabilities with lossless - rejection sampling. - -3. **Commit** - - The emitted step is `accepted drafts + bonus token`. - - The cached prefix is `anchor + accepted drafts`. - - The bonus token is emitted but not cached, because it becomes the next - anchor. - -4. **Rollback** - - Serving keeps `cached_len == prompt_len + emitted_len` after each committed - step. That is the state contract pinned by `SpecDecodeState`. - - `SpecDecodeState.record_verify(block_len, committed)` advances the cached - and emitted counts, and `apply_verify_step()` turns that into the - `truncate_target` passed to `PagedKVCache.truncate_tokens(seq_id, new_len)`. - - `PagedKVCache.truncate_tokens(seq_id, new_len)` frees tail blocks and - truncates swapped-out buffers. - - Sliding-window targets snapshot and replay the committed prefix when plain - `DynamicCache.crop()` is not enough. - -## Sampled decoding - -- Sampled DFlash is batch-1 only. -- The direct path uses `warped_probs`, `acceptance_sampled`, and - `committed_tokens_sampled`. -- Warp order matches the engine sampler, temperature first, then top-k, then - top-p, then softmax. -- If you use `MoE.generate` or `MoE.serve`, sampled requests stay on the - standard path today, except that Qwen3.5-MoE intentionally rejects sampled - speculative decode with `ValueError` on the `MoE.generate` path. -- Batch > 1 sampled DFlash is not supported. - -## Batched decoding - -- Batch > 1 direct DFlash is greedy only and requires a bare HuggingFace target. -- Prompts must be left padded, or all the same length. `attention_mask` must be - 0/1 valued and end with a real token. -- `max_new_tokens` can be a single int or a per-sequence list. -- The returned tensor is right padded, and `last_generated_lengths` stores the - true per-row new-token counts. -- `MoE.generate(..., speculative_draft=...)` with batch > 1 raises - `NotImplementedError`. - -## Continuous-batching serving - -- `MoE.serve(..., speculative_draft=...)` and the OpenAI server CLI accept a - DFlash speculator. -- The serving engine delegates only for a fresh singleton prefill request with - no prior output tokens that is greedy, stop-free, penalty-free, has - `top_k <= 0`, `top_p >= 1.0`, `repetition_penalty == 1.0`, `logprobs <= 0`, - and is within the per-step token cap. -- A delegated serving step may emit several accepted tokens, and the engine - streams them one by one after recording committed counts. -- The validated real-checkpoint serving harness is GPT-OSS-20B. GPT-OSS-120B - has no recorded real-checkpoint serving validation. Qwen3.5 serving wiring - has tiny-fixture coverage, but no real-checkpoint serving validation is - recorded. - -## Route-ahead expert prefetch - -Route-ahead is a scheduling-only add-on. It warms the exact routed expert union -for the layer being dispatched, but it does not change routing or output -tokens. - -- Activation happens only during a DFlash verify forward. -- If route-ahead is active and a prefetcher exists, the executor pins the exact - routed union with `fetch_experts_lock_cache(...)` and enqueues the same set - with `speculative_prefetch(..., expert_ids=..., prefetch_layer_id=...)`. -- If the context is inactive, there is no prefetcher, or the union is empty, - the executor falls back to the legacy pooled prefetch or a no-op. -- The union is pinned one layer at a time because `ReplaceCacheCandidates` is - global and clears background queues. -- Offloaded executor-backed models such as DeepSeek, Qwen, and Mixtral reach - this seam. GPT-OSS does not, because its expert path never wires an - `expert_executor`. - -Metric names: - -- `RouteAheadStepSummary(layers, predicted, actual, covered, kept, wasted)` -- `RouteAheadStats.as_dict()` returns `steps`, `layers_observed`, - `predicted_experts`, `actual_experts`, `covered_experts`, `kept_experts`, - `wasted_experts`, `coverage`, and `waste_ratio` - -Coverage is `covered / actual` and defaults to `1.0` when `actual == 0`. -Waste ratio is `wasted / predicted` and defaults to `0.0` when -`predicted == 0`. - -## Observability - -- `route_ahead_stats` defaults to `None`. `enable_route_ahead_stats()` returns - the recorder and resets it on reuse. -- `step_trace`, `last_target_cache`, `last_draft_cache`, and - `last_generated_lengths` are advanced internal diagnostics only; they have no - stability or compatibility guarantee. -- `step_trace` records `prev_start`, `accept`, `start`, `emitted_len`, - `target_cache_len`, and `draft_cache_len`. -- `last_target_cache` and `last_draft_cache` expose the final caches after a run. -- `last_generated_lengths` is only set by the batched path. - -## Compatibility - -| Model | Target / drafter | Residency | Sync | Serving | Sampling | Route-ahead | Hardware / validation | -| --- | --- | --- | --- | --- | --- | --- | --- | -| GPT-OSS-20B | `openai/gpt-oss-20b` / `z-lab/gpt-oss-20b-DFlash` | Tunable via `offload_path` and `device_memory_ratio`; the GPU-gated harness uses a high resident setting | Greedy batch-1 validated | Continuous-batching greedy validated | Direct batch-1 sampled implemented, not yet validated on the real pair | Not wired: the GPT-OSS expert path does not attach an `expert_executor` | GPU-gated, no board asserted in repo tests | -| GPT-OSS-120B | `openai/gpt-oss-120b` / `z-lab/gpt-oss-120b-DFlash` | Tunable via `offload_path` and `device_memory_ratio`; the GPU-gated harness uses a high resident setting | Greedy batch-1 validated | Implemented, not yet validated on a real serving harness | Direct batch-1 sampled implemented, not yet validated on the real pair | Not wired: the GPT-OSS expert path does not attach an `expert_executor` | GPU-gated, no board or TP count asserted in repo tests | -| Qwen3.5-MoE | No published real-model drafter checkpoint is recorded in repo tests | Text backbone, shared expert, and `lm_head` stay resident, routed experts offload | Greedy batch-1 hybrid rollback validated on tiny CPU fixtures | Serving wiring has tiny-fixture coverage; no real-checkpoint serving validation is recorded | `MoE.generate` sampled DFlash is rejected, direct sampled speculator is not yet validated on a real checkpoint | Wired and exercised on synthetic/tiny CPU fixtures; not validated with a real target/drafter checkpoint | No real-model hardware validation recorded | - -Route-ahead is an execution-path capability of executor-backed offloaded models, -not evidence that a validated DFlash target/drafter checkpoint pair exists. -DeepSeek, Qwen, and Mixtral paths can reach the executor seam described above, -but they are absent from this matrix unless the repo records a corresponding -DFlash pairing and validation scope. No additional drafter pair is implied. - -The real-model harnesses skip cleanly unless `MOE_DFLASH_GPU=1` is set and the -checkpoints are present in the HuggingFace cache. - -## Validation - -- CPU gate, no GPU or checkpoint download required: - `pytest -q tests/python/dflash -m "not gpu"` -- GPT-OSS-20B GPU-gated validation: - `MOE_DFLASH_GPU=1 pytest -q tests/python/dflash/test_gpu_20b_dflash.py` -- GPT-OSS-120B GPU-gated validation: - `MOE_DFLASH_GPU=1 pytest -q tests/python/dflash/test_gpu_120b.py` -- GPT-OSS-20B serving versus sync validation: - `MOE_DFLASH_GPU=1 pytest -q tests/python/dflash/test_gpu_serving_dflash.py` - -## Known limitations and troubleshooting - -- Batch > 1 with `speculative_draft` through `MoE.generate` is not supported. -- Batch > 1 sampled DFlash is not supported. -- GPT-OSS route-ahead is not wired because the model path never attaches an - `expert_executor`. -- Qwen3.5 greedy DFlash is supported through the hybrid cache replay path, but - sampled DFlash on a real Qwen3.5 checkpoint is not yet validated. -- Missing checkpoints or `MOE_DFLASH_GPU` simply skip the GPU-gated harnesses. -- `block_size`, `target_layer_ids`, hidden size, or vocab mismatches fail fast - during drafter validation. - -## Related files - -- `examples/dflash_gpt_oss_example.py` -- `tests/python/dflash/` +## Evidence boundaries + +- Published GPT-OSS pairs: pairing/direct evidence; no executor route-ahead. +- DeepSeek V2/V3: executor and default-off MLA paging capability; no real + DeepSeek DFlash pair claim. +- Qwen3.5: tiny-fixture direct/hybrid evidence only; no real-pair claim. +- Hybrid paged rollback and GPT-OSS paged MLA are not implemented claims. diff --git a/docs/environment-variables.md b/docs/environment-variables.md index 273ed2ba..311a8da3 100644 --- a/docs/environment-variables.md +++ b/docs/environment-variables.md @@ -53,7 +53,9 @@ For `CUDA_VISIBLE_DEVICES` ordering, expert ownership, and one-host multi-GPU be | `NVTX_DISABLE` | `"0"` | `setup.py` build | If `"1"`, compile out NVTX instrumentation macros. | Build-time only. | | `MOE_ENABLE_SM90` | `"1"` | `setup.py` build | Include sm_90 kernels in the compiled extensions. | Build-time only. | | `MOE_ENABLE_SM120` | `"0"` | `setup.py` build | Include sm_120 kernels and the native FP4 extension arch flags. | Build-time only. | -| `MOEINF_VERSION` | `"0.0.1"` | `setup.py` packaging | Set the package version string. | Packaging only. | + +The package version is not set via an environment variable; it is derived from +git tags at build time by setuptools-scm (see `pyproject.toml`). ## Standard third-party envs diff --git a/docs/model-compatibility.md b/docs/model-compatibility.md index f325cedb..2c0e226c 100644 --- a/docs/model-compatibility.md +++ b/docs/model-compatibility.md @@ -1,42 +1,72 @@ # Model compatibility matrix -This guide is the source of truth for the families registered in -`moe_infinity/common/constants.py` and the runtime adapters wired in -`moe_infinity/runtime/model_offload.py`. - -Legend: -- `validated`, covered by a repo test on a real checkpoint or a close smoke harness -- `implemented/experimental`, the code path exists, but the repo only has unit or fixture coverage -- `not validated`, the repo has no harness for the claim -- `unsupported`, no runtime path or a fail-fast guard blocks it - -Conditional registry keys, `deepseekv4`, `qwen3_5`, and `glmmoedsa`, only -register when the matching Transformers class imports cleanly. The repo floor -is `transformers>=5.3.0,<6`, but some families need newer builds. - -| Family / HF class | Example checkpoint | Minimum Transformers | Sync generation | Continuous serving | Expert offload / quantization | Speculative decoding | Validated topology | Limitations | -|---|---|---|---|---|---|---|---|---| -| DeepSeek-V2 (`DeepseekV2ForCausalLM`) | `deepseek-ai/DeepSeek-V2-Lite-Chat` | `>= 5.3.0` | validated | implemented/experimental | validated | not recorded | `1x GPU` | FlashAttention is excluded in the offload path, eager attention is used in the consistency harness | -| DeepSeek-V3 (`DeepseekV3ForCausalLM`) | `deepseek-ai/DeepSeek-V3` | `>= 5.3.0` | implemented/experimental | implemented/experimental | implemented/experimental | not recorded | Not recorded | Only routing and paged-attention parity are covered in repo tests | -| DeepSeek-V4 (`DeepseekV4ForCausalLM`) | `deepseek-ai/DeepSeek-V4-Flash` | Not recorded, guarded import | validated for the official offload path | not validated | validated | unsupported | `4x GPU mp4` | The official checkpoint must be mp-sharded; repo validation covers the mp4 path, while mp1 is not covered by repo tests | -| Mixtral (`MixtralForCausalLM`) | `mistralai/Mixtral-8x7B-Instruct-v0.1` | `>= 5.3.0` | implemented/experimental | implemented/experimental | implemented/experimental | not recorded | Not recorded | No real-model harness is recorded in this repo | -| Qwen3 (`Qwen3MoeForCausalLM`) | `Qwen/Qwen3-30B-A3B` | `>= 5.3.0` | validated | implemented/experimental | validated | not recorded | `1x GPU` | `Qwen3PagedAttention` exists, but the repo does not ship a real-model serving harness for it | -| Qwen3.5 (`Qwen3_5MoeForConditionalGeneration`) | `Qwen/Qwen3.5-35B-A3B` | `>= 5.12` | validated on tiny fixtures through deprecated `MoE.generate()` | serving wiring covered by tiny fixtures; no real-checkpoint serving validation recorded | validated on tiny fixtures | validated on tiny fixtures | Tiny CPU fixtures only | Text-only path, vision and MTP weights stay unused, batch>1 speculative draft is blocked in deprecated `MoE.generate()` | -| GLM-5.2 (`GlmMoeDsaForCausalLM`) | `zai-org/GLM-5.2-FP8` | `>= 5.12` | validated | validated on the tiny serving harness | validated | built-in MTP only | `1x GPU` | Native engine is forced off, no GLM DFlash drafter is registered, and non-zero temperature falls back to greedy in MTP | -| GPT-OSS (`GptOssForCausalLM`) | `openai/gpt-oss-20b` | `>= 5.3.0` | validated for 20B | validated for 20B | validated | validated for greedy batch-1, sampled batch-1 implemented/experimental | GPU-gated | Route-ahead is not wired because the path never attaches an `expert_executor`; batch>1 `speculative_draft` is blocked | -| DBRX (`DbrxForCausalLM`) | `databricks/dbrx-instruct` | `>= 5.3.0` | implemented/experimental | not validated | implemented/experimental | not recorded | Not recorded | Registry and adapter code exist, but there is no real-model harness in repo tests | -| Jamba (`JambaForCausalLM`) | `ai21labs/Jamba-*` | `>= 5.3.0` | implemented/experimental | not validated | implemented/experimental | not recorded | Not recorded | Registry and adapter code exist, but there is no real-model harness in repo tests | -| OLMoE (`OlmoeForCausalLM`) | `allenai/OLMoE-*` | `>= 5.3.0` | implemented/experimental | not validated | implemented/experimental | not recorded | Not recorded | Registry and adapter code exist, but there is no real-model harness in repo tests | -| NLLB-MoE (`NllbMoeForConditionalGeneration`) | `facebook/nllb-moe-54b` | `>= 5.3.0` | implemented/experimental | not validated | implemented/experimental | not recorded | Not recorded | Encoder-decoder sparse parsing is covered, but there is no repo end-to-end harness | -| OPT (`OPTForCausalLM`) | Not recorded | `>= 5.3.0` | unsupported | unsupported | unsupported | unsupported | Not recorded | Registry entry only, no adapter or test coverage beyond parser dispatch | - -Repo evidence: -- `moe_infinity/common/constants.py` for the registry and conditional imports -- `moe_infinity/runtime/model_offload.py` for the adapter wiring -- `tests/test_gpt_oss_*`, `tests/python/v4/*`, `tests/python/integration/test_glm_*`, - `tests/python/unit/test_qwen3_5_moe.py`, `tests/python/unit/test_model_registry.py`, - `tests/python/unit/test_glm_*`, and `tests/python/integration/test_model_consistency.py` - -Use `Not recorded` when the repo has no direct evidence for a version, -topology, or capability claim. Use `implemented/experimental` when the code path -exists but only unit or fixture tests cover it. +This page separates general model support from DFlash evidence. `validated` +means the named scope has a repository harness; `implemented/experimental` +means code plus tiny/unit evidence; `not recorded` means no direct evidence. +Pairing validity never implies executor reachability, and a rich or paged claim +is valid only behind its corresponding capability declaration. + +Loading some published drafter checkpoints requires +`trust_remote_code=True`. This executes arbitrary code supplied by the remote +repository. Only load a trusted, pinned drafter revision, and review that +revision before use. This warning applies to pairing examples below; it is not +a claim that any particular DeepSeek DFlash pair has been validated. + +## General model support + +| Family / HF class | General sync/offload status | Continuous serving | Notes | +| --- | --- | --- | --- | +| DeepSeek-V2 (`DeepseekV2ForCausalLM`) | validated | implemented/experimental | Eager consistency harness; Stage 4b details below. | +| DeepSeek-V3 (`DeepseekV3ForCausalLM`) | implemented/experimental | implemented/experimental | Routing and paged-attention parity evidence. | +| DeepSeek-V4 (`DeepseekV4ForCausalLM`) | validated official mp4 path | not validated | DFlash unsupported; mp1 not covered. | +| Mixtral (`MixtralForCausalLM`) | implemented/experimental | implemented/experimental | No real-model serving harness recorded. | +| Qwen3 / Qwen3.5 MoE | Qwen3 validated; Qwen3.5 tiny-fixture validated | implemented/experimental | Qwen3.5 is text-only and requires newer Transformers. | +| GLM-5.2 (`GlmMoeDsaForCausalLM`) | validated | tiny serving harness | Built-in MTP, not DFlash. | +| GPT-OSS (`GptOssForCausalLM`) | 20B validated | 20B validated | Resident expert implementation. | +| DBRX / Jamba / OLMoE / NLLB-MoE | implemented/experimental | not validated | Registry/adapter evidence only. | +| OPT (`OPTForCausalLM`) | unsupported | unsupported | Registry entry only. | + +## DFlash capability and evidence matrix + +| Family | DFlash pairing evidence | Direct execution | Rich execution capability | Executor / route-ahead evidence | Serving cache capability | Validation boundary | +| --- | --- | --- | --- | --- | --- | --- | +| GPT-OSS-20B | valid published pairs: `openai/gpt-oss-20b` / `z-lab/gpt-oss-20b-DFlash` | Greedy real-pair GPU evidence; sampled direct implementation with tiny statistical evidence | Resident wrapper; no paged MLA claim | **no executor route-ahead** because the resident path does not attach `expert_executor` | Stage 4a compatibility path; no GPT-OSS paged MLA | GPU fixture required for real-pair claims. | +| GPT-OSS-120B | valid published pairs: `openai/gpt-oss-120b` / `z-lab/gpt-oss-120b-DFlash` | Greedy GPU-gated harness; sampled direct implementation not real-pair validated | Resident wrapper | **no executor route-ahead** | Stage 4a compatibility path | No board/TP claim is inferred. | +| DeepSeek V2/V3 | **No real DeepSeek DFlash pair** is recorded | Session semantics and local/tiny adapter evidence only | MLA rows use grouped per-request fallback; physical rich batching needs the row-aware capability guard | Executor seam is reachable on offloaded models; this does not validate pairing | `paged_mla` is default-off, eligible only for batch-1 greedy DeepSeek V2/V3; exactly one engine-owned target paged store per request, using `MLAPagedKVCache` for packed MLA; drafter cache separate; resident-only and no swap/preemption | DeepSeek MLA uses the correct PyTorch fallback, not FlashInfer acceleration. Stage 4b proves ownership, not a checkpoint pair. | +| Qwen3.5-MoE | No published pair recorded | Qwen tiny-only evidence; greedy hybrid rollback fixtures | **Qwen/hybrid fallback** is grouped per request, not physically batched | Executor seam has synthetic/tiny evidence | Stage 4a `temporary_dynamic`; no hybrid paged rollback | No real checkpoint/drafter or sampled serving claim. | +| Other executor-backed MoE | Not recorded | Not recorded | Physical rich batching only after the wrapper's row-aware capability guard | Wiring may exist | No DFlash paged claim | Executor evidence is not pairing evidence. | + +### Capability gates + +- **Rich execution capability:** `supports_batch` and `supports_rich_forward` + must both be true, and the wrapper must preserve row-aligned logits, hidden + states, cache handles, masks, positions, and route contexts. Otherwise + scheduling may group rows but forwards remain per request. +- **Serving cache capability:** Stage 4b requires the default-off DeepSeek MLA + flag, a compatible MLA module set, target-cache-adapter support, greedy mode, + and batch 1. Every other path uses the explicit Stage 4a fallback. +- **Sampling:** direct bare-HF batch 1/batch > 1 supports greedy, sampled, and + mixed rows with per-row RNG. This does not widen deprecated or paged-serving + surfaces automatically. +- **Paged-store ownership:** Each eligible request has exactly one engine-owned + target paged store, either standard `PagedKVCache` or packed-MLA + `MLAPagedKVCache`. The drafter cache is separate. DRAFT/VERIFY speculative + sessions are resident and non-preemptible. The implemented active-session cap + and post-peak free-block reserve (declared budget plus transient verify + headroom) route rejected admissions immediately to + Stage 4a; they do not prove general fairness. No swap/resume claim is made. +- **Preemption:** ordinary serving sequences can swap/preempt. Paged MLA DFlash + is resident-only; no swap/preemption while DRAFT/VERIFY is in flight. + +## Evidence sources + +- `tests/python/dflash/test_capability_orthogonality.py` +- `tests/python/dflash/test_bare_hf_backend.py` +- `tests/python/dflash/test_mixed_sampling_batch.py` +- `tests/python/dflash/test_rich_batch_forward.py` +- `tests/python/serving/test_dflash_stage4a.py` +- `tests/python/serving/test_dflash_stage4b.py` +- `tests/python/serving/test_rich_batch_runner.py` + +Use `Not recorded` rather than extrapolating hardware, pairing, route-ahead, +sampling, rich batching, or paged-cache support from an adjacent capability. diff --git a/docs/serving.md b/docs/serving.md index 1a3aa9bf..3bbed6e9 100644 --- a/docs/serving.md +++ b/docs/serving.md @@ -167,24 +167,63 @@ python -m moe_infinity.entrypoints.openai.api_server_v2 \ --speculative-draft z-lab/gpt-oss-20b-DFlash ``` -Startup validates the drafter/target pair: hidden size, vocab size, mask-token bounds, target layer IDs, and drafter `fc` shape. - -Delegation is server-wide and only applies when a request is: - -- a fresh singleton prefill request -- greedy (`temperature=0`, no sampling) -- `top_k <= 0` -- `top_p >= 1.0` -- `repetition_penalty == 1.0` -- `logprobs <= 0` -- no stop strings -- within the current step token budget - -The exact gate is implemented in [`moe_infinity/serving/engine.py`](../moe_infinity/serving/engine.py) and also requires batch==1, no prior output tokens, and `max_tokens <= scheduler.max_tokens_per_step`. - -The delegated path runs the speculative loop and emits tokens normally through SSE. - -Route-ahead is internal to the DFlash verify path and is used only when speculative delegation is active. +Startup validates structural pairing (hidden size, vocabulary, mask-token +bounds, target layers, block constraints, and drafter shape) separately from +executor/route-ahead reachability. + +The persistent path creates one canonical session per eligible sequence. It +preserves the request's temperature, top-k, top-p, budget, EOS set, and +request-scoped generator. It does not silently turn sampled requests into +greedy requests. Unsupported grammar/guided/logit-bias metadata, penalties, +logprobs, or stop strings use the standard serving fallback before drafting. +That fallback is not evidence of sampled serving. + +Two cache execution contexts are observable in `/admin/stats`: + +- `temporary_dynamic` is the Stage 4a compatibility mode. It keeps a temporary + private DynamicCache while the engine owns scheduling, callbacks, + cancellation, and request accounting. Sampled sessions and ineligible model + layouts remain here; this is not sampled paged-MLA serving. +- `paged_mla` is the Stage 4b default-off target enabled by + `enable_deepseek_mla_paging=True`. It is restricted to eligible greedy + batch-1 DeepSeek V2/V3 MLA sessions. The engine owns packed latent/rope target + pages; the draft cache is separate. Admission is bounded by + `max_resident_paged_speculative_sessions` (default `1`) and must leave at + least `min_free_mla_blocks_after_admission` free blocks (default `1`) after + reserving the block-rounded peak for the full declared + `prompt + max_tokens` budget plus up to `DFlash block_size - 1` transient + verify tokens. Active sessions' committed and transient headroom that is not + yet allocated is included. A rejected eligible request immediately uses + `temporary_dynamic`; it does not wait for a paged seat, and its sampling + parameters are unchanged. That dense Stage 4a fallback owns a private target + cache and can therefore increase total GPU memory use even though it consumes + no MLA pages. + +Paged MLA is currently resident-only and has no preemption/swap implementation. +The scheduler does not preempt DRAFT/VERIFY sessions. Qwen and hybrid layouts +fall back to Stage 4a; hybrid paged rollback is not claimed. Cancellation after +an in-flight backend call releases session resources, and per-sequence page +ownership prevents one cancellation or rollback from truncating another row. +Completion/cancellation frees ownership, so a later request can be admitted. +`/admin/stats` reports active paged sessions, current free blocks, configured +limits, and counters for `admitted`, `session_cap`, `free_block_reserve`, and +`ineligible` decisions. `begin_failed` is recorded only when adapter/session +construction fails; `admitted` increments only after construction succeeds. + +The guard is block-based admission control, not a general fairness proof. It +does not preempt or swap admitted sessions. External cache consumers can still +invalidate reserved headroom; such allocator failures clean up the affected +request and are currently re-raised by the engine step. + +There is no real DeepSeek DFlash target/drafter pair validation in the repo. +Stage 4b's tiny/local DeepSeek adapter tests establish ownership and attention +metadata only. GPT-OSS has named valid pairs, but its resident expert path has +no executor route-ahead. Qwen evidence is tiny-fixture only. + +Route-ahead is observer-only. Pairing evidence, executor reachability, +prefetch-fired evidence, and cache ownership are reported as separate facts. +See [DFlash unified execution](dflash.md) for direct batching, RNG caveats, +benchmarks, and exact CPU/GPU gates. ## Operational Endpoints diff --git a/docs/superpowers/plans/2026-08-17-dflash-unified-execution.md b/docs/superpowers/plans/2026-08-17-dflash-unified-execution.md new file mode 100644 index 00000000..b9322594 --- /dev/null +++ b/docs/superpowers/plans/2026-08-17-dflash-unified-execution.md @@ -0,0 +1,177 @@ +# DFlash Unified Execution Implementation Plan + +**Goal:** Deliver one DFlash semantic core with capability-selected direct, +deprecated-sync, rich, and serving execution. +**Status:** Tasks 1-12 implemented and reviewed; Task 13 records the actual +delivered behavior and rollout gates. +**Corrected dependency order:** Task 8 -> Task 8.5 -> Task 9 -> Tasks 10-13. + +This checked-in plan is the canonical plan revised after implementation. It +records departures from the initial ordering without rewriting them as if they +had always been known. + +## Non-negotiable guards + +1. Sampling claims require retained proposal rows and request RNG. +2. Grouped rich requests are not physical rich model batching. +3. Pairing, rich execution, executor reachability, and prefetch firing are + separate evidence dimensions. +4. Stage 4b cannot be selected without a model-owned representation that writes + exactly one engine-owned target paged store per eligible request. The store + is standard `PagedKVCache` or packed-MLA `MLAPagedKVCache`; the drafter cache + is separate. +5. Missing required GPU fixtures are not successful rollout gates. + +## Task record + +### Task 1: Shared protocol and trace schema — delivered + +Added request/sampling/result/capability/cache contracts plus `SessionTrace`, +`PairingEvidence`, and `ExecutorEvidence`. + +### Task 2: Request-scoped RNG — delivered + +Threaded generators through anchor, drafter, acceptance, and correction draws; +retained every sampled proposal row through verification and reconstruction. + +### Task 3: Single direct session — delivered + +Moved batch-1 direct semantics onto canonical sessions while retaining output, +cache, stop, and trace behavior. + +### Task 4: Session driver and cohorts — delivered + +Added capability-first backend selection, safe cohort splitting, fail-before- +output behavior, progress checks, cleanup, and atomic physical-cohort results. + +### Task 5: Bare-HF greedy physical batching — delivered + +Moved left-padding, position IDs, row commit, dense rollback/reconstruction, +right-padding, and `last_generated_lengths` into the bare-HF backend. + +### Task 6: Sampled and mixed bare-HF batching — delivered + +Enabled batch-1/batch>1 sampled and mixed rows with per-row policies and RNG. +One scalar generator is cloned and can correlate rows; cross-batch-shape +bit-exactness is not guaranteed. + +### Task 7: Public/deprecated facade migration — delivered with compatibility limits + +Direct APIs share session semantics. `MoE.generate()` retains its +`DeprecationWarning`, return rectangle, and greedy compatibility behavior. It +does not claim general sampled batch > 1 support; Qwen3.5 keeps its explicit +non-greedy rejection. + +### Task 8: Stage 4a serving — delivered + +Added persistent per-sequence canonical sessions with `temporary_dynamic` +execution, verify demand, streaming commits, failure records, cancellation, and +logical invariant checks. + +### Task 8.5: DeepSeek MLA prerequisite — inserted and delivered before Task 9 + +**Why reordered:** Task 9's page adapter could not establish final target-cache +ownership while the DeepSeek model still produced ordinary dense KV. The +DeepSeek MLA prerequisite restored/adapted the paging foundation from history: + +- default-off DeepSeek V2/V3 eligibility; +- engine-owned packed latent/rope `MLAPagedKVCache`; +- per-layer attention adaptation and ownership validation; +- rich-forward attention metadata and position handling; +- tiny/local parity and ownership tests. + +The corrected dependency is **Task 8.5 -> Task 9**. Task 9 was blocked until +this foundation existed. + +### Task 9: Stage 4b paged ownership — delivered in constrained scope + +Added `PagedCacheAdapter` and `paged_mla` execution for default-off eligible +batch-1 greedy DeepSeek V2/V3. Target pages are engine-owned; the draft cache is +separate. Cancellation and completion release only the owning sequence. + +The actual delivered behavior is narrower than the initial final-state wording: + +- all DRAFT/VERIFY speculative sessions are resident and non-preemptible; +- Stage 4a temporarily double-allocates target state; +- resident-only execution trades GPU memory for progress; the implemented + default-one active-session cap and block-rounded declared request-budget plus + transient verify-peak reserve reject to Stage 4a immediately with observable + reasons/counters; dense fallback may increase total GPU memory; +- release permits later admission, but the guard is not a general fairness + proof against unrelated cache consumers; +- there is no speculative swap/resume claim; +- sampled and ineligible paths retain Stage 4a; +- Qwen/hybrid paths retain Stage 4a; +- there is no hybrid paged rollback claim; +- DeepSeek MLA uses the correct PyTorch fallback, not FlashInfer acceleration; +- there is no real DeepSeek DFlash target/drafter pair claim. + +### Task 10: Pairing/executor orthogonality — delivered + +Separated pairing evidence from executor evidence and made route-ahead failure +observer-only. GPT-OSS keeps valid named pairs but no executor route-ahead. + +### Task 11: Independent loop retirement — delivered + +`_generate_batched` now normalizes compatibility arguments and calls the +physical session driver/backend; it no longer owns a second acceptance loop. + +### Task 12: Rich row awareness — delivered with capability gate + +Added row metadata/results and physical rich execution for wrappers explicitly +declaring the full contract. Tiny standard-cache fixtures validate physical +batching. MLA and hybrid/Qwen wrappers fall back to grouped per-request +execution. Qwen evidence remains tiny-only. + +### Task 13: Documentation, matrix, benchmark, and rollout gates — delivered here + +Files: + +- `docs/dflash.md`, `docs/model-compatibility.md`, `docs/serving.md` +- `README.md`, `ARCHITECTURE.md`, `CHANGELOG.md` +- `tests/python/dflash/test_compatibility_matrix.py` +- `benchmarks/dflash/unified_execution_benchmark.py` +- `benchmarks/dflash/validate_unified_execution.py` + +The compatibility assertions gate rich and paged claims on explicit capability +language and keep pairing/executor columns separate. The tiny benchmark performs +real CPU protocol timing and sampled-law measurements; it labels its synthetic +scope and does not fabricate checkpoint metrics. The validator fails closed on +cache invariant, ownership, order, sampling, or required-GPU failures. + +## Validation commands + +No-download CPU gates: + +```bash +pytest -q tests/python/dflash/test_compatibility_matrix.py +python benchmarks/dflash/unified_execution_benchmark.py --fixture tiny +python benchmarks/dflash/validate_unified_execution.py --fixture tiny \ + --require-cache-invariants --require-order-invariance +pytest -q tests/python/dflash +pytest -q tests/python/serving +``` + +GPU gates: + +```bash +MOE_DFLASH_GPU=1 CUDA_VISIBLE_DEVICES=0 \ + pytest -q tests/python/dflash/test_gpu_20b_dflash.py \ + tests/python/dflash/test_gpu_serving_dflash.py -m gpu +MOE_DFLASH_GPU=1 CUDA_VISIBLE_DEVICES=0 \ + pytest -q tests/python/dflash/test_gpu_120b.py -m gpu +python benchmarks/dflash/validate_unified_execution.py --fixture tiny \ + --require-cache-invariants --require-order-invariance --require-gpu +``` + +`--require-gpu is a readiness gate`: the validator returns failure unless CUDA +and `MOE_DFLASH_GPU=1` are both present, but it does not execute the GPU harness. +The actual GPU pytest command remains separate and required. A skipped +checkpoint test is reported as unavailable, not passed. + +## Completion boundary + +Task 13 completes truthful documentation and local rollout validation. It does +not remove Stage 4a, enable Stage 4b by default, establish sampled paged serving, +claim a DeepSeek/Qwen real pair, claim GPT-OSS route-ahead, claim hybrid paged +rollback, or convert unavailable GPU evidence into success. diff --git a/docs/superpowers/plans/2026-08-18-dflash-loop-retirement.md b/docs/superpowers/plans/2026-08-18-dflash-loop-retirement.md new file mode 100644 index 00000000..d43d0354 --- /dev/null +++ b/docs/superpowers/plans/2026-08-18-dflash-loop-retirement.md @@ -0,0 +1,111 @@ +# DFlash Legacy Batch Loop Retirement Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Make `_generate_batched` a compatibility adapter whose requests execute only through a capability-selected `SessionDriver` physical cohort. + +**Architecture:** Add a physical-cohort run result and `SessionDriver.run_physical_cohort` beside the existing request-scoped `run` lifecycle. The driver validates every request and selects one compatible batch backend before model execution, calls the backend once, validates the complete returned cohort, then atomically publishes row-level `DriverResult` objects and traces. The legacy tensor method retains its signature and handles only normalization, request construction, driver invocation, diagnostics, and rectangular output adaptation. + +**Tech Stack:** Python, PyTorch, pytest, Python `ast`, basedpyright/LSP. + +--- + +### Task 1: Lock Down Loop Retirement and Driver Semantics + +**Files:** +- Create: `tests/python/dflash/test_loop_retirement.py` +- Test: `tests/python/dflash/test_loop_retirement.py` + +- [ ] **Step 1: Write failing structural tests** + +Parse `DFlashSpeculator._generate_batched` with `ast`, resolve called attribute/name tails, and assert it contains no decode `while` loop and no calls to acceptance, committed-token, warped-probability, target/draft forward, cache snapshot/rollback, or route-stat operations. Assert that it constructs `RequestSpec`, `BatchedBareHFBackend`, and `SessionDriver`, and invokes `run_physical_cohort`. + +- [ ] **Step 2: Write failing physical-cohort driver tests** + +Use a fake `PhysicalCohortBackend` to prove all `supports`/capability checks precede `execute_cohort`, sampled requests never downgrade to a non-sampling backend, unsupported or incompatible rows fail before execution, returned row count/budgets/traces are validated atomically, successful rows become ordered `DriverResult` values, and backend exceptions leave `last_results` empty. + +- [ ] **Step 3: Write failing adapter parity tests** + +Cover greedy, sampled, and mixed requests with scalar/per-row budgets, sampling contexts/generators, stop sets, and masks. Verify one physical driver invocation, right-padded output, batch-one shape, `last_generated_lengths`, caches, aggregate step trace, row `DriverResult` values, session traces, and route/pairing/executor evidence. + +- [ ] **Step 4: Verify RED** + +Run: `pytest -q tests/python/dflash/test_loop_retirement.py` + +Expected: failures because `SessionDriver.run_physical_cohort` and the physical run result do not exist and `_generate_batched` still invokes `execute_cohort` directly. + +### Task 2: Add the Physical Cohort Driver Entry + +**Files:** +- Modify: `moe_infinity/spec_decode/session_driver.py` +- Modify: `moe_infinity/spec_decode/__init__.py` +- Test: `tests/python/dflash/test_loop_retirement.py` + +- [ ] **Step 1: Define the physical run result** + +Add an immutable `PhysicalCohortDriverResult` containing ordered `DriverResult` rows plus the opaque backend cohort result so compatibility adapters can retain backend-owned cache diagnostics without duplicating execution logic. + +- [ ] **Step 2: Implement capability-first physical selection** + +Normalize and uniquely identify requests, filter to runtime `PhysicalCohortBackend` implementations with `supports_batch`, preserve sampling capability checks, validate hashable/common cohort keys, and reject a physical cohort that would require backend/key splitting before calling any backend execution method. + +- [ ] **Step 3: Implement one atomic physical execution** + +Validate tensor/mask rank and row counts, derive shared or per-row stops and row sampling contexts, call `execute_cohort` once, validate generated rows and session traces for every request, convert them to `DriverResult` values, and only then assign `last_results` and return the wrapper. On every exception retain `last_results == ()`. + +- [ ] **Step 4: Verify focused driver tests pass** + +Run: `pytest -q tests/python/dflash/test_loop_retirement.py -k 'driver or physical'` + +Expected: PASS. + +### Task 3: Reduce `_generate_batched` to a Legacy Adapter + +**Files:** +- Modify: `moe_infinity/spec_decode/dflash.py` +- Update: `tests/python/dflash/test_bare_hf_backend.py` +- Test: `tests/python/dflash/test_loop_retirement.py` + +- [ ] **Step 1: Construct request rows from legacy arguments** + +Keep the existing method signature. Validate shape/budgets/mask without model execution, strip left padding for each `RequestSpec.prompt_token_ids`, preserve each `SamplingContext`, budget, and stop set, and retain legacy shared-stop resolution. + +- [ ] **Step 2: Invoke the driver and adapt the complete result** + +Construct `BatchedBareHFBackend(self)` and `SessionDriver([backend])`, call `run_physical_cohort`, copy backend cache/step diagnostics and driver row results/traces only after success, compute generated lengths, right-pad new tokens with the target pad id, and concatenate with the original padded prompt tensor. + +- [ ] **Step 3: Verify adapter and structural tests pass** + +Run: `pytest -q tests/python/dflash/test_loop_retirement.py tests/python/dflash/test_bare_hf_backend.py` + +Expected: PASS. + +### Task 4: Regression and Static Verification + +**Files:** +- Verify: `moe_infinity/spec_decode/session_driver.py` +- Verify: `moe_infinity/spec_decode/dflash.py` +- Verify: `moe_infinity/spec_decode/__init__.py` +- Verify: `tests/python/dflash/test_loop_retirement.py` + +- [ ] **Step 1: Run all DFlash tests** + +Run: `pytest -q tests/python/dflash` + +Expected: PASS, with only environment-declared skips. + +- [ ] **Step 2: Run relevant public/serving regression tests** + +Run: `pytest -q tests/python/dflash/test_public_api_compat.py tests/python/dflash/test_engine_wire.py tests/python/serving/test_dflash_stage4a.py tests/python/serving/test_dflash_stage4b.py` + +Expected: PASS, with only environment-declared skips. + +- [ ] **Step 3: Run static diagnostics** + +Run LSP diagnostics on every changed Python file, then run the repository's configured basedpyright command over the changed implementation and test files. + +Expected: no errors. + +- [ ] **Step 4: Confirm scope and worktree state** + +Inspect the diff and status without staging or committing. Confirm no serving/MLA files were changed by Task11 and stop after the first complete successful verification. diff --git a/docs/superpowers/plans/2026-08-18-paged-mla-admission-guard.md b/docs/superpowers/plans/2026-08-18-paged-mla-admission-guard.md new file mode 100644 index 00000000..8f440b7a --- /dev/null +++ b/docs/superpowers/plans/2026-08-18-paged-mla-admission-guard.md @@ -0,0 +1,114 @@ +# Paged MLA Admission Guard Implementation Plan + +> **For agentic workers:** REQUIRED SUB-SKILL: Use superpowers:subagent-driven-development (recommended) or superpowers:executing-plans to implement this plan task-by-task. Steps use checkbox (`- [ ]`) syntax for tracking. + +**Goal:** Bound resident, non-preemptible `paged_mla` sessions and preserve an MLA free-block reserve by immediately routing rejected admissions to the existing Stage 4a path. + +**Architecture:** `ArcherConfig` and serving-engine config expose validated block-based limits. `MLAPagedKVCache` exposes a read-only free-block count, while `SpecSessionDriver` makes an atomic synchronous admission decision before allocation using the block-rounded declared budget plus transient verify peak, records a structured decision, and maintains counters surfaced by engine stats. Existing session release naturally removes the active-session count and makes later admissions eligible. + +**Tech Stack:** Python dataclasses, PyTorch cache allocator, pytest. + +--- + +### Task 1: Configuration and cache introspection + +**Files:** +- Modify: `moe_infinity/utils/config.py` +- Modify: `moe_infinity/serving/mla_cache.py` +- Test: `tests/python/unit/test_utils_config.py` +- Test: `tests/python/serving/test_mla_paged_cache.py` + +- [ ] Add failing tests asserting defaults of one resident session and one reserved free MLA block, rejection of booleans/negative session caps/reserves below one, and a public `free_block_count` that tracks allocation and release. +- [ ] Run `pytest -q tests/python/unit/test_utils_config.py tests/python/serving/test_mla_paged_cache.py` and confirm the new assertions fail for missing fields/API. +- [ ] Add `max_resident_paged_speculative_sessions: int = 1` and `min_free_mla_blocks_after_admission: int = 1`, validate exact integer types with cap `>= 0` and reserve `>= 1`, and expose `MLAPagedKVCache.free_block_count` by reading the allocator count. + +### Task 2: Admission, fallback, lifecycle, and statistics + +**Files:** +- Modify: `moe_infinity/serving/spec_session_driver.py` +- Modify: `moe_infinity/serving/engine.py` +- Test: `tests/python/serving/test_dflash_stage4b.py` + +- [ ] Add failing tests for the concurrent-session cap, insufficient post-allocation reserve, Stage 4a fallback without sampling changes, admission after release, structured per-record decisions/counters, and default-off unchanged behavior. +- [ ] Run `pytest -q tests/python/serving/test_dflash_stage4b.py` and confirm failures identify the missing admission policy. +- [ ] Pass validated engine limits into `SpecSessionDriver`; before creating `PagedCacheAdapter`, count live `paged_mla` records and compute peak demand as `ceil((prompt_tokens + max_new_tokens + dflash_block_size - 1) / cache.block_size)`. Admit only when below the cap and `free_block_count - demand >= reserve`; otherwise create the existing temporary-dynamic context immediately. +- [ ] Store reason codes (`admitted`, `session_cap`, `free_block_reserve`, or `ineligible`) on records, increment admission counters, and expose policy, active count, free blocks, and counters from `ContinuousBatchingEngine.get_stats()`. +- [ ] Run `pytest -q tests/python/serving/test_dflash_stage4b.py tests/python/serving/test_dflash_stage4a.py` and confirm all pass. + +### Task 3: Serving configuration and documentation + +**Files:** +- Modify: `moe_infinity/entrypoints/openai/api_server_v2.py` +- Modify: `docs/serving.md` +- Modify: `docs/dflash.md` +- Modify: `docs/superpowers/specs/2026-08-17-dflash-unified-execution-design.md` +- Modify: `docs/superpowers/plans/2026-08-17-dflash-unified-execution.md` + +- [ ] Add serving CLI/config defaults for both guard fields and pass them into the engine config. +- [ ] Replace “guard required” wording with the implemented cap/reserve/fallback behavior and retain limitations: resident/no swap, no preemption, block-based estimate, no general fairness proof. + +### Task 4: Verification + +**Files:** All modified Python files. + +- [ ] Run language-server diagnostics on every modified Python file and require zero errors. +- [ ] Run `pytest -q tests/python/unit/test_utils_config.py tests/python/serving/test_mla_paged_cache.py tests/python/serving/test_dflash_stage4b.py tests/python/serving/test_dflash_stage4a.py` once; stop after the first successful verification. + +### Follow-up: Full-budget reservation and serving configuration + +**Files:** +- Modify: `moe_infinity/serving/spec_session_driver.py` +- Modify: `moe_infinity/entrypoints/openai/api_server_v2.py` +- Modify: `moe_infinity/entrypoints/big_modeling.py` +- Test: `tests/python/serving/test_dflash_stage4b.py` +- Test: `tests/python/serving/test_api_routes.py` + +- [ ] Add failing exact-fit and one-block-short tests using + `ceil((prompt_len + max_new_tokens + dflash_block_size - 1) / mla_block_size)`, including active + sessions' declared-but-not-yet-allocated headroom. +- [ ] Add a failing backend-begin test proving `admitted` is not incremented + until adapter retention and session construction succeed, while + `begin_failed` is incremented and prompt pages are released. +- [ ] Implement per-session block budgets and dynamically subtract active + unallocated headroom before deciding whether the candidate plus configured + reserve fits. +- [ ] Add failing CLI/default/config-builder/programmatic-initializer tests for + `enable_deepseek_mla_paging`, + `max_resident_paged_speculative_sessions`, and + `min_free_mla_blocks_after_admission`. +- [ ] Forward the three fields through `parse_args`, `_build_engine_config`, CLI + `moe_config`, `initialize_with_model`, and `MoE.serve`, preserving defaults + `False`, `1`, and `1`. +- [ ] Add an engine `add_request` -> `step` paged-selection test and a forced + external MLA exhaustion test proving draft failure releases the request's + pages, records `speculative_draft_failed`, and re-raises the allocator error. +- [ ] Update serving/DFlash docs to state that declared prompt-plus-output + capacity is reserved and Stage 4a dense fallback can increase total GPU + memory. +- [ ] Run the requested admission/API/serving/DFlash tests and language-server + diagnostics once after implementation. + +### Final polish: transient verify peak + +**Files:** +- Modify: `moe_infinity/serving/spec_session_driver.py` +- Modify: `moe_infinity/utils/config.py` +- Modify: `moe_infinity/entrypoints/openai/api_server_v2.py` +- Test: `tests/python/serving/test_dflash_stage4b.py` +- Test: `tests/python/serving/test_api_routes.py` + +- [ ] Add failing exact-fit and one-block-short tests where the DFlash verify + block is larger than the MLA cache block. Reserve + `ceil((prompt + max_new_tokens + dflash_block_size - 1) / mla_block_size)`. +- [ ] Add a test that appends the declared committed budget plus maximum + transient verify tokens without allocator exhaustion after admission. +- [ ] Change standalone `SpecSessionDriver` paging to default off and make every + direct paged test opt in explicitly; keep the engine's explicit flag wiring. +- [ ] Store transient-inclusive peak block budgets for active sessions so their + unallocated headroom remains reserved. +- [ ] Update CLI/config help and docs to describe full declared capacity plus + transient verify headroom, then run Ruff, LSP, and focused tests. +- [ ] Cover missing, non-integer, zero, and one-valued DFlash block sizes as + ineligible Stage 4a fallbacks with no MLA allocation. Help text defines the + reserve as blocks remaining after all active/new declared budgets and + transient verify peaks. diff --git a/docs/superpowers/specs/2026-08-17-dflash-unified-execution-design.md b/docs/superpowers/specs/2026-08-17-dflash-unified-execution-design.md new file mode 100644 index 00000000..fc96f77d --- /dev/null +++ b/docs/superpowers/specs/2026-08-17-dflash-unified-execution-design.md @@ -0,0 +1,141 @@ +# DFlash Unified Execution Design + +**Date:** 2026-08-17 +**Status:** Approved design, updated with delivered scope on 2026-08-18 +**Dependency order:** Task 8 -> Task 8.5 -> Task 9 -> Tasks 10-13 + +## Decision + +DFlash uses one request-oriented semantic core for draft, verify, acceptance, +sampling, stop handling, rollback, and traces. Model execution is selected by +explicit capabilities. Direct generation and the deprecated sync facade use +`SessionDriver`; serving preserves the same `SpecSession` transitions through a +lifecycle-owning driver. `_generate_batched` is an adapter rather than a second +semantic loop. + +## Semantic contracts + +Each request owns its prompt, budget, stop set, anchor, proposal rows, target and +draft cache handles, request RNG, emitted tokens, and trace. Sampled anchor, +draft, acceptance, and correction draws consume the request stream. No sampled +decision may use another row's generator or silently become greedy. + +The cache contract after commit is: + +```text +target cache = prompt + committed non-bonus prefix +anchor = emitted bonus, not duplicated in target cache +``` + +Rejected verify-tail cache is truncated before another round. Dense batches may +re-feed already chosen tokens only to reconstruct a shared physical rectangle; +proposal rows and RNG state remain authoritative. + +## Backends and capabilities + +### Bare Hugging Face + +The physical bare-HF backend supports batch 1 and batch > 1, greedy, sampled, +and mixed rows. Prompts are left padded, outputs are right padded, and +`last_generated_lengths` carries true generated lengths. Scalar generators are +cloned per row and can correlate identical requests; explicit per-row generators +are recommended. Cross-batch-shape bit-exactness is not promised. + +### Rich MoE + +A per-request rich backend is the safe baseline. Physical rich batching requires +row-aligned input metadata, logits, hidden states, cache handles, masks, +positions, route contexts, and executor row unions. Wrappers declare this +capability explicitly. Grouped scheduling is not physical batching. MLA and +hybrid/Qwen wrappers retain per-request fallback unless their exact cache +contract is declared. + +### Serving Stage 4a + +`temporary_dynamic` is a compatibility context. It uses private dense target +state while `ContinuousBatchingEngine` retains scheduler, callback, +cancellation, and accounting ownership. It is temporary architecture, not +paged-cache ownership. + +### Task 8.5: DeepSeek MLA prerequisite + +The original ordering placed Task 9 before the model could write an engine-owned +DeepSeek MLA representation. Task 8.5 was inserted as the **DeepSeek MLA +prerequisite** and restored/adapted the historical paging foundation before +Stage 4b. It supplies packed `[kv_c_normed | k_pe]` storage, attention metadata, +DeepSeek V2/V3 eligibility, model adaptation, and the engine-owned target-cache +handle. The required order is **Task 8.5 -> Task 9**. + +### Serving Stage 4b + +Stage 4b is default-off. `enable_deepseek_mla_paging=True` may select +`paged_mla` only for eligible batch-1 greedy DeepSeek V2/V3 MLA requests. Each +eligible request has exactly one engine-owned target paged store, either the +standard `PagedKVCache` or packed-MLA `MLAPagedKVCache`; the drafter cache +remains separate. All in-flight DRAFT/VERIFY speculative sessions are resident +and non-preemptible. Stage 4a temporarily double-allocates target state, and +resident-only execution trades GPU memory for progress. The implemented guard +defaults to one active paged session and one free MLA block after expected +full declared prompt-plus-output allocation plus up to one DFlash verify +block's transient tail (`block_size - 1` tokens), including active sessions' +unallocated peak headroom; cap/reserve rejection records a reason and immediately uses +Stage 4a rather than waiting. Released sessions open admission for later +requests. Dense Stage 4a fallback may increase total GPU memory. This is bounded +admission, not a general fairness proof against unrelated cache consumers. +There is no swap/resume claim. Sampled, +Qwen, hybrid, and other ineligible paths remain Stage 4a. DeepSeek MLA +currently uses the correct PyTorch fallback rather than FlashInfer acceleration. +This design does not claim a real DeepSeek DFlash checkpoint pair or hybrid +paged rollback. + +## Evidence dimensions + +Pairing and execution are orthogonal. `PairingEvidence` covers DFlash config, +dimensions, vocabulary, mask, layers, block constraints, module validation, and +named checkpoint scope. `ExecutorEvidence` covers executor reachability, +attempted/fired route layers, actual expert unions, bytes, coverage, and +fallback. GPT-OSS named pairs can be valid while route-ahead is unreachable; +executor-wired DeepSeek/Qwen/Mixtral models can lack a validated pair. + +Route-ahead is observer-only. Missing wiring or prefetch failure cannot modify +routing, acceptance, outputs, or cache state. + +## Serving lifecycle + +The engine owns admission, scheduling, verify demand, callbacks, cancellation, +and completion. Cancellation waits for an in-flight backend call and then +releases the session. Standard non-speculative sequences retain scheduler +preemption. In-flight DRAFT/VERIFY sessions are resident and non-preemptible. +Per-sequence ownership prevents cross-request truncation, but no speculative +swap/resume behavior is claimed. Admission never cancels or preempts an +existing session. + +## Observability + +The shared logical trace includes request ID, backend, cache kind, sampled mode, +rounds, accepted, committed, emitted, rollback, replay, finish reason, pairing +evidence, and executor evidence. Rich metrics separate per-request from physical +calls. Benchmarks must label fixture scope and report unavailable capabilities +as unavailable, never as zero-valued successes. +Serving stats additionally report the paged-MLA policy, active/free counts, and +decision counters. + +## Delivered compatibility boundary + +| Surface | Delivered behavior | +| --- | --- | +| Direct bare HF | Greedy/sampled/mixed, batch 1 and batch > 1 | +| Direct rich | Batch 1; grouped per-request fallback; physical tiny-fixture batching behind capability | +| Deprecated `MoE.generate` | Warning-preserving compatibility facade; no general sampled batch > 1 widening | +| Serving Stage 4a | Persistent canonical sessions with temporary dynamic target state | +| Serving Stage 4b | Default-off eligible greedy batch-1 DeepSeek V2/V3 MLA only | +| GPT-OSS | Named pair evidence; no executor route-ahead | +| DeepSeek | MLA paging/ownership evidence; no real DFlash pair claim | +| Qwen | Tiny/hybrid evidence only; no real pair or hybrid paged rollback | + +## Non-goals + +This design does not claim sampled paged serving, a real DeepSeek pair, +GPT-OSS paged MLA or executor route-ahead, hybrid paged rollback, universal rich +batching, GPU success from skipped tests, or release status from an unreleased +changelog entry. diff --git a/moe_infinity/__init__.py b/moe_infinity/__init__.py index b173f8d0..44033e61 100644 --- a/moe_infinity/__init__.py +++ b/moe_infinity/__init__.py @@ -1,4 +1,9 @@ -__version__ = "0.0.1" +try: + # Generated at build time by setuptools-scm (see pyproject.toml). + from ._version import __version__ +except Exception: # pragma: no cover - source tree built without setuptools-scm + __version__ = "0.0.0+unknown" + __all__ = ["MoE", "OffloadEngine", "__version__"] diff --git a/moe_infinity/distributed/expert_executor.py b/moe_infinity/distributed/expert_executor.py index 37072b7d..63feca1b 100644 --- a/moe_infinity/distributed/expert_executor.py +++ b/moe_infinity/distributed/expert_executor.py @@ -94,6 +94,27 @@ def _layer_expert_nbytes(prefetcher, layer_id, expert_ids): return entry or None +def _executor_evidence(**kwargs): + # Lazy for the same package-cycle reason as ``_load_route_ahead_impl``. + from moe_infinity.spec_decode.protocols import ExecutorEvidence + + return ExecutorEvidence(**kwargs) + + +def _prefetcher_hit_rate(prefetcher): + getter = getattr(prefetcher, "get_hit_rate", None) + if not callable(getter): + return None + try: + value = getter() + except Exception: + return None + if isinstance(value, bool) or not isinstance(value, (int, float)): + return None + rate = float(value) + return rate if 0.0 <= rate <= 1.0 else None + + class DistributedExpertExecutor: def __init__(self, archer_config: ArcherConfig): self.archer_config = archer_config @@ -104,6 +125,11 @@ def __init__(self, archer_config: ArcherConfig): getattr(archer_config, "speculative_prefetch_overlap", False) ) self._pending_prefetch = None + self._pending_prefetch_failure_safe = False + self.last_executor_evidence = _executor_evidence( + wiring_reachable=True, + fallback_reason="context_inactive", + ) def set_expert_dispatcher(self, expert_dispatcher): global _expert_dispatcher @@ -144,6 +170,12 @@ def _maybe_route_ahead_prefetch( """ ctx, union_experts_from_mask = _load_route_ahead_impl() if not ctx.is_active(): + available_prefetcher = prefetcher or self.prefetcher + self.last_executor_evidence = _executor_evidence( + wiring_reachable=True, + prefetcher_present=available_prefetcher is not None, + fallback_reason="context_inactive", + ) return False stats = ctx.current_stats() route_prefetcher = prefetcher @@ -153,7 +185,22 @@ def _maybe_route_ahead_prefetch( route_prefetcher = self.prefetcher mask_2d = router_mask.reshape(-1, num_expert) union_expert_ids = union_experts_from_mask(mask_2d) + row_union: set[tuple[int, int, int]] = set() + row_offsets = ctx.current_row_offsets() + if row_offsets and row_offsets[-1] == int(mask_2d.shape[0]): + for row in range(len(row_offsets) - 1): + row_ids = union_experts_from_mask( + mask_2d[row_offsets[row] : row_offsets[row + 1]] + ) + row_union.update( + (row, int(layer_id), int(expert_id)) + for expert_id in row_ids + ) fired = False + fallback_reason = None + expert_nbytes = _layer_expert_nbytes( + route_prefetcher, layer_id, union_expert_ids + ) if route_prefetcher is not None and union_expert_ids: # A0 section 2/5 (A4 guard): pin exactly ONE layer's union per # dispatch -- ``ReplaceCacheCandidates`` is global and clears the @@ -161,27 +208,76 @@ def _maybe_route_ahead_prefetch( # into one pin would evict candidates the next layer's dispatch # still needs. Never batch pins across layers; never pin the # empty set (short-circuited above). - route_prefetcher.fetch_experts_lock_cache( - layer_id, union_expert_ids - ) - route_prefetcher.speculative_prefetch( - layer_id, - expert_ids=union_expert_ids, - prefetch_layer_id=layer_id, - ) - fired = True + try: + route_prefetcher.fetch_experts_lock_cache( + layer_id, union_expert_ids + ) + route_prefetcher.speculative_prefetch( + layer_id, + expert_ids=union_expert_ids, + prefetch_layer_id=layer_id, + ) + fired = True + except Exception as exc: + # Route-ahead is cache-warming observation only. A prefetch + # failure must preserve the legacy expert dispatch/output path. + fallback_reason = f"prefetch_exception:{type(exc).__name__}" + elif not union_expert_ids: + fallback_reason = "empty_actual_union" + else: + fallback_reason = "prefetcher_absent" + + prefetched_bytes = ( + sum(expert_nbytes.values()) + if fired and expert_nbytes is not None + else 0 + ) + cache_hit_rate = _prefetcher_hit_rate(route_prefetcher) + self.last_executor_evidence = _executor_evidence( + wiring_reachable=True, + prefetcher_present=route_prefetcher is not None, + attempted_layers=(int(layer_id),), + fired_layers=((int(layer_id),) if fired else ()), + actual_expert_union=frozenset( + (int(layer_id), int(expert_id)) + for expert_id in union_expert_ids + ), + actual_expert_union_by_row=frozenset(row_union), + prefetched_bytes=prefetched_bytes, + coverage=(1.0 if fired or not union_expert_ids else 0.0), + cache_hit_rate=cache_hit_rate, + fallback_reason=fallback_reason, + ) if stats is not None: # A5 read-only observation: predicted == the pinned union when # the prefetch fired, else [] (coverage 0 for this layer). predicted_ids = union_expert_ids if fired else [] - stats.observe_layer( - layer_id, - predicted_ids, - mask_2d, - expert_nbytes=_layer_expert_nbytes( - route_prefetcher, layer_id, predicted_ids - ), - ) + observe_attempt = getattr(stats, "observe_executor_attempt", None) + if callable(observe_attempt): + try: + observe_attempt( + layer_id, + union_expert_ids, + actual_ids_by_row=row_union, + prefetcher_present=route_prefetcher is not None, + fired=fired, + fallback_reason=fallback_reason, + prefetched_bytes=prefetched_bytes, + cache_hit_rate=cache_hit_rate, + ) + except Exception: + # Observer failures are isolated from expert dispatch. + pass + try: + stats.observe_layer( + layer_id, + predicted_ids, + mask_2d, + expert_nbytes=(expert_nbytes if fired else None), + ) + except Exception: + # Observer failures are isolated from expert dispatch. + pass return fired def dispatch_local( @@ -225,6 +321,9 @@ def dispatch_local( route_ahead_handled = self._maybe_route_ahead_prefetch( layer_id, router_mask, num_expert, prefetcher ) + route_ahead_attempted = bool( + self.last_executor_evidence.attempted_layers + ) dispatch_nvtx_ctx = _nvtx_ctx("expert_dispatch") dispatch_profiler_ctx = ( @@ -255,7 +354,13 @@ def dispatch_local( and prefetcher is not None and router_logits is not None ): - self.trigger_speculative_prefetch(layer_id, router_logits) + if route_ahead_attempted: + try: + self.trigger_speculative_prefetch(layer_id, router_logits) + except Exception: + pass + else: + self.trigger_speculative_prefetch(layer_id, router_logits) pending_router_logits = None else: pending_router_logits = router_logits @@ -266,6 +371,7 @@ def dispatch_local( expert_list, pending_router_logits, ) + self._pending_prefetch_failure_safe = route_ahead_attempted def wait_dispatch_local(self): profiler = _profiler_instance() @@ -283,10 +389,26 @@ def wait_dispatch_local(self): if pending is not None: prefetcher, layer_id, expert_list, router_logits = pending self._pending_prefetch = None + failure_safe = self._pending_prefetch_failure_safe + self._pending_prefetch_failure_safe = False if prefetcher is not None: - prefetcher.correct_prefetch(layer_id + 1, expert_list) + if failure_safe: + try: + prefetcher.correct_prefetch(layer_id + 1, expert_list) + except Exception: + pass + else: + prefetcher.correct_prefetch(layer_id + 1, expert_list) if router_logits is not None: - self.trigger_speculative_prefetch(layer_id, router_logits) + if failure_safe: + try: + self.trigger_speculative_prefetch( + layer_id, router_logits + ) + except Exception: + pass + else: + self.trigger_speculative_prefetch(layer_id, router_logits) return result diff --git a/moe_infinity/engine/generation_loop.py b/moe_infinity/engine/generation_loop.py index a3c9b765..218c4025 100644 --- a/moe_infinity/engine/generation_loop.py +++ b/moe_infinity/engine/generation_loop.py @@ -114,6 +114,9 @@ def _spec_strategy_applies( ) -> bool: if self.spec_strategy is None: return False + supports = getattr(self.spec_strategy, "supports_engine_request", None) + if callable(supports): + return bool(supports(sp, batch_size=batch_size)) if batch_size != 1: return False if ( diff --git a/moe_infinity/entrypoints/big_modeling.py b/moe_infinity/entrypoints/big_modeling.py index 951f78f5..23b4683d 100644 --- a/moe_infinity/entrypoints/big_modeling.py +++ b/moe_infinity/entrypoints/big_modeling.py @@ -250,6 +250,21 @@ def __init__( is_flash_attn_available=is_flash_attn_available, trust_remote_code=True, ) + mla_cache = native_components.get("mla_cache") + if mla_cache is not None: + from moe_infinity.models.deepseek_mla_attention import ( + adapt_deepseek_model, + ) + + adapted = adapt_deepseek_model( + self.model, + mla_cache, + enabled=True, + ) + if not adapted: + self._native_mla_cache = None + + self._native_rich_batch_capable = self._supports_native_rich_batch() self._ensure_generation_mixin() self.engine_config = engine_config @@ -314,6 +329,7 @@ def _build_native_components( self._native_memory_coordinator = None self._native_kv_cache_manager = None self._native_attention_backend = None + self._native_mla_cache = None self._native_transfer_scheduler = None self._native_scheduler = None self._native_generation_engine = None @@ -411,20 +427,46 @@ def _build_native_components( ) num_cpu_blocks = max(32, num_gpu_blocks * 2) + from moe_infinity.models.deepseek_mla_attention import ( + is_deepseek_mla_eligible, + ) + + mla_enabled = is_deepseek_mla_eligible( + model_config, + enabled=bool( + getattr(engine_config, "enable_deepseek_mla_paging", False) + ), + ) + kv_cache_manager = KVCacheManager( num_gpu_blocks=num_gpu_blocks, num_cpu_blocks=num_cpu_blocks, block_size=kv_spec.block_size, ) device = torch.device("cuda" if torch.cuda.is_available() else "cpu") - try: - attention_backend = PagedAttentionBackend( - spec=kv_spec, - num_gpu_blocks=num_gpu_blocks, + mla_cache = None + if mla_enabled: + from moe_infinity.serving.mla_cache import MLAPagedKVCache + + mla_cache = MLAPagedKVCache( + num_blocks=num_gpu_blocks, + block_size=kv_spec.block_size, + num_layers=int(num_layers), + latent_dim=int(getattr(model_config, "kv_lora_rank")), + rope_dim=int(getattr(model_config, "qk_rope_head_dim")), + dtype=kv_spec.dtype, device=device, ) - except Exception: attention_backend = None + else: + try: + attention_backend = PagedAttentionBackend( + spec=kv_spec, + num_gpu_blocks=num_gpu_blocks, + device=device, + ) + except Exception: + attention_backend = None transfer_scheduler = UnifiedTransferScheduler() kv_offload_coordinator = None if getattr(engine_config, "enable_kv_cache_offload", False): @@ -493,6 +535,7 @@ def _build_native_components( self._native_memory_coordinator = memory_coordinator self._native_kv_cache_manager = kv_cache_manager self._native_attention_backend = attention_backend + self._native_mla_cache = mla_cache self._native_transfer_scheduler = transfer_scheduler self._native_scheduler = scheduler self._native_generation_engine = generation_engine @@ -503,6 +546,7 @@ def _build_native_components( "memory_coordinator": memory_coordinator, "kv_cache_manager": kv_cache_manager, "attention_backend": attention_backend, + "mla_cache": mla_cache, "transfer_scheduler": transfer_scheduler, "kv_offload_coordinator": kv_offload_coordinator, "expert_offload_coordinator": expert_offload_coordinator, @@ -583,10 +627,14 @@ def _native_model_forward( if not is_prefill and _attention_metadata is not None: seq_lens = getattr(_attention_metadata, "seq_lens", None) if seq_lens is not None and seq_lens.numel() > 0: - current_pos = int(seq_lens[0].item()) - 1 - extra_kwargs["position_ids"] = torch.tensor( - [[current_pos]], device=input_tensor.device - ) + total_len = int(seq_lens[0].item()) + start_pos = total_len - len(token_ids) + extra_kwargs["position_ids"] = torch.arange( + start_pos, + total_len, + device=input_tensor.device, + dtype=torch.long, + ).unsqueeze(0) with torch.no_grad(): if not use_paged_context: @@ -622,10 +670,10 @@ def _native_model_forward( def _native_model_forward_rich( self, - token_ids: list[int], + token_ids: list[int] | torch.Tensor, _attention_metadata: object = None, logits_to_keep: int = 0, - ) -> tuple[torch.Tensor, tuple, object]: + ) -> tuple[torch.Tensor, tuple, object] | object: """On-device forward for speculative decoding: hidden-state capture. Single HF forward with ``output_hidden_states=True`` returning @@ -642,16 +690,40 @@ def _native_model_forward_rich( ExpertExecutor dispatch (and its ``speculative_prefetch`` hook) is preserved — nothing here bypasses expert dispatch. """ - input_tensor = torch.tensor([token_ids], dtype=torch.long) + batched = isinstance(token_ids, torch.Tensor) + if batched: + if token_ids.ndim != 2: + raise ValueError( + "batched rich token_ids must have shape [batch, seq]" + ) + input_tensor = token_ids.to(dtype=torch.long) + else: + input_tensor = torch.tensor([token_ids], dtype=torch.long) input_tensor = input_tensor.to(self._resolve_native_input_device()) is_prefill = True if _attention_metadata is not None: is_prefill = bool(getattr(_attention_metadata, "is_prefill", True)) + mla_attention_modules = self._get_mla_attention_modules() + use_mla_context = bool( + mla_attention_modules + and _attention_metadata is not None + and getattr(self, "_native_mla_cache", None) is not None + and all( + hasattr(_attention_metadata, name) + for name in ( + "block_tables", + "seq_lens", + "slot_mapping", + "is_prefill", + ) + ) + ) paged_attention_classes = self._get_paged_attention_classes() use_paged_context = bool( - paged_attention_classes + not use_mla_context + and paged_attention_classes and _attention_metadata is not None and getattr(self, "_native_attention_backend", None) is not None ) @@ -659,27 +731,63 @@ def _native_model_forward_rich( extra_kwargs: dict = {"output_hidden_states": True} if logits_to_keep: extra_kwargs["logits_to_keep"] = int(logits_to_keep) + if batched and _attention_metadata is not None: + attention_mask = getattr( + _attention_metadata, "attention_mask", None + ) + position_ids = getattr(_attention_metadata, "position_ids", None) + if attention_mask is not None: + extra_kwargs["attention_mask"] = attention_mask.to( + input_tensor.device + ) + if position_ids is not None: + extra_kwargs["position_ids"] = position_ids.to( + input_tensor.device + ) - if not use_paged_context: + if not use_paged_context and not use_mla_context: # Same HF KV-cache contract as the baseline: prefill captures # past_key_values; decode steps consume the cached KV. extra_kwargs["use_cache"] = True - if not is_prefill: - cached_kv = getattr(self, "_cached_past_key_values", None) + if not is_prefill or batched: + cached_kv = None + handles = getattr(_attention_metadata, "cache_handles", ()) + if handles and all(handle is handles[0] for handle in handles): + cached_kv = handles[0] + if cached_kv is None: + cached_kv = getattr(self, "_cached_past_key_values", None) if cached_kv is not None: extra_kwargs["past_key_values"] = cached_kv else: + extra_kwargs["use_cache"] = False if not is_prefill and _attention_metadata is not None: seq_lens = getattr(_attention_metadata, "seq_lens", None) if seq_lens is not None and seq_lens.numel() > 0: - current_pos = int(seq_lens[0].item()) - 1 - extra_kwargs["position_ids"] = torch.tensor( - [[current_pos]], device=input_tensor.device - ) + total_len = int(seq_lens[0].item()) + start_pos = total_len - len(token_ids) + extra_kwargs["position_ids"] = torch.arange( + start_pos, + total_len, + device=input_tensor.device, + dtype=torch.long, + ).unsqueeze(0) with torch.no_grad(): - if not use_paged_context: + if not use_paged_context and not use_mla_context: outputs = self.model(input_tensor, **extra_kwargs) + elif use_mla_context: + from moe_infinity.models.deepseek_mla_attention import ( + clear_deepseek_mla_context, + set_deepseek_mla_context, + ) + + for module in mla_attention_modules: + set_deepseek_mla_context(module, _attention_metadata) + try: + outputs = self.model(input_tensor, **extra_kwargs) + finally: + for module in mla_attention_modules: + clear_deepseek_mla_context(module) else: backend = self._native_attention_backend for attn_cls in paged_attention_classes: @@ -690,7 +798,7 @@ def _native_model_forward_rich( for attn_cls in paged_attention_classes: attn_cls.clear_paged_context() - if not use_paged_context: + if not use_paged_context and not use_mla_context: past_kv = getattr(outputs, "past_key_values", None) if past_kv is not None: self._cached_past_key_values = past_kv @@ -710,9 +818,48 @@ def _native_model_forward_rich( "output_hidden_states=True is required" ) + if use_mla_context or batched: + from moe_infinity.spec_decode.protocols import RichForwardResult + + cache_handle = ( + self._native_mla_cache + if use_mla_context + else getattr(outputs, "past_key_values", None) + ) + supplied_handles = tuple( + getattr(_attention_metadata, "cache_handles", ()) + ) + if use_mla_context and supplied_handles: + row_handles = supplied_handles + else: + row_handles = (cache_handle,) * int(input_tensor.shape[0]) + + return RichForwardResult( + logits=logits, + hidden_states=tuple(hidden_states), + cache_handle=cache_handle, + cache_handles=row_handles, + row_offsets=tuple( + getattr(_attention_metadata, "row_offsets", ()) + ), + row_lengths=tuple( + getattr(_attention_metadata, "row_lengths", ()) + ), + ) past_key_values = getattr(outputs, "past_key_values", None) return logits, hidden_states, past_key_values + def _get_mla_attention_modules(self) -> list[torch.nn.Module]: + names = {"DeepseekV2MLAPagedAttention", "DeepseekV3MLAPagedAttention"} + modules_fn = getattr(self.model, "modules", None) + if not callable(modules_fn): + return [] + return [ + module + for module in modules_fn() + if module.__class__.__name__ in names + ] + def _get_paged_attention_classes(self) -> list[type[Any]]: paged_class_names = { "DeepseekV2PagedAttention", @@ -740,6 +887,26 @@ def _get_paged_attention_classes(self) -> list[type[Any]]: return classes + def _supports_native_rich_batch(self) -> bool: + """Fail-closed declaration for the dense row-aware rich contract.""" + config = getattr(self.model, "config", None) + if any( + bool(getattr(config, name, False)) + for name in ( + "hybrid_attention", + "sliding_window_pattern", + "recurrent_chunk_size", + ) + ): + return False + # Existing paged wrappers retain engine ownership and are driven by + # ModelRunner; MLA and Qwen/hybrid rollback are intentionally not + # widened through the dense direct-generation backend. + return ( + not self._get_paged_attention_classes() + and not self._get_mla_attention_modules() + ) + def _configure_hook(self, input_ids: torch.LongTensor): if self.arch == "mixtral": import moe_infinity.models.mixtral # noqa: F401 @@ -827,19 +994,30 @@ def generate(self, input_ids: torch.LongTensor, **kwargs) -> Any: speculative_draft = kwargs.pop("speculative_draft", None) + generation_config = kwargs.get("generation_config") + + def generation_value(name: str, default: Any) -> Any: + if name in kwargs: + return kwargs[name] + if generation_config is not None: + value = getattr(generation_config, name, None) + if value is not None: + return value + return default + model_type = getattr( getattr(self.model, "config", None), "model_type", "" ) is_qwen35 = model_type == "qwen3_5_moe" - do_sample = kwargs.get("do_sample", None) + do_sample = generation_value("do_sample", None) sampling_temperature = ( - 0.0 if do_sample is False else float(kwargs.get("temperature", 1.0)) - ) - is_greedy = ( - sampling_temperature == 0.0 - and float(kwargs.get("top_p", 1.0)) == 1.0 - and int(kwargs.get("top_k", 0)) == 0 + 0.0 + if do_sample is False + else float(generation_value("temperature", 1.0)) ) + top_p = float(generation_value("top_p", 1.0)) + top_k = int(generation_value("top_k", 0) or 0) + is_greedy = sampling_temperature == 0.0 and top_p == 1.0 and top_k == 0 qwen35_dflash = bool(speculative_draft) and is_greedy native_engine = self._native_generation_engine native_for_call = ( @@ -850,6 +1028,88 @@ def generate(self, input_ids: torch.LongTensor, **kwargs) -> Any: and (not is_qwen35 or qwen35_dflash) ) + if speculative_draft and is_qwen35 and not is_greedy: + raise ValueError( + "qwen3_5_moe speculative_draft (DFlash) requires " + "greedy decoding (do_sample=False or temperature=0, " + "top_p=1, top_k=0)" + ) + + if ( + speculative_draft + and is_greedy + and input_ids.ndim == 2 + and input_ids.shape[0] > 1 + ): + if not self.use_native_engine or native_engine is None: + raise ValueError( + "speculative_draft (DFlash) requires the MoE-Infinity " + "native engine (use_native_engine=True)" + ) + self._resolve_spec_strategy(speculative_draft) + try: + speculator = getattr(self, "_dflash_speculator", None) + if speculator is None: + raise RuntimeError("DFlash speculator was not configured") + + attention_mask = kwargs.get("attention_mask") + if attention_mask is None: + prompt_lengths = [int(input_ids.shape[1])] * int( + input_ids.shape[0] + ) + else: + if tuple(attention_mask.shape) != tuple(input_ids.shape): + raise ValueError( + f"attention_mask shape {tuple(attention_mask.shape)} != " + f"input_ids shape {tuple(input_ids.shape)}" + ) + prompt_lengths = [ + int(value) + for value in attention_mask.to(dtype=torch.long) + .sum(dim=1) + .tolist() + ] + for prompt_length in prompt_lengths: + if prompt_length > self.max_seq_length: + raise ValueError( + f"prompt length {prompt_length} exceeds max_seq_length " + f"{self.max_seq_length}" + ) + + max_tokens = generation_value("max_new_tokens", None) + if max_tokens is None: + max_tokens = generation_value("max_tokens", 256) + stop_token_ids = kwargs.get("stop_token_ids") + if stop_token_ids is None: + stop_token_ids = generation_value("eos_token_id", None) + if isinstance(stop_token_ids, bool): + raise ValueError( + "eos_token_id/stop_token_ids cannot be boolean" + ) + if isinstance(stop_token_ids, int): + stop_token_ids = [stop_token_ids] + generator = kwargs.get("generator") + self._configure_hook(input_ids) + self._cached_past_key_values = None + self.model.eval() + output = speculator.generate( + input_ids, + max_new_tokens=max_tokens, + temperature=sampling_temperature, + stop_token_ids=stop_token_ids, + top_k=top_k, + top_p=top_p, + attention_mask=attention_mask, + generator=generator, + ) + self.last_dflash_traces = getattr( + speculator, "last_session_traces", () + ) + return output + finally: + self._cached_past_key_values = None + native_engine.spec_strategy = None + if not native_for_call: if speculative_draft: if input_ids.ndim == 2 and input_ids.shape[0] != 1: @@ -857,12 +1117,6 @@ def generate(self, input_ids: torch.LongTensor, **kwargs) -> Any: "speculative_draft (DFlash) v1 supports batch==1 " f"only; got batch size {input_ids.shape[0]}" ) - if is_qwen35 and not is_greedy: - raise ValueError( - "qwen3_5_moe speculative_draft (DFlash) requires " - "greedy decoding (do_sample=False or temperature=0, " - "top_p=1, top_k=0)" - ) raise ValueError( "speculative_draft (DFlash) requires the MoE-Infinity " "native engine (use_native_engine=True)" @@ -890,11 +1144,13 @@ def generate(self, input_ids: torch.LongTensor, **kwargs) -> Any: f"prompt length {len(prompt_token_ids)} exceeds max_seq_length {self.max_seq_length}" ) - max_tokens = kwargs.get("max_new_tokens", kwargs.get("max_tokens", 256)) + max_tokens = generation_value("max_new_tokens", None) + if max_tokens is None: + max_tokens = generation_value("max_tokens", 256) sampling_params = SamplingParams( temperature=sampling_temperature, - top_p=float(kwargs.get("top_p", 1.0)), - top_k=int(kwargs.get("top_k", 0)), + top_p=top_p, + top_k=top_k, max_tokens=int(max_tokens) if max_tokens is not None else 256, ) try: @@ -922,6 +1178,9 @@ def serve( enable_prefix_caching: bool = False, offload_dir: Optional[str] = None, speculative_draft: Optional[object] = None, + enable_deepseek_mla_paging: bool = False, + max_resident_paged_speculative_sessions: int = 1, + min_free_mla_blocks_after_admission: int = 1, ) -> None: """ Start the OpenAI-compatible continuous batching server. @@ -936,6 +1195,9 @@ def serve( kv_cache_ratio: Fraction of device_memory_ratio for KV cache (default: 0.25) max_batch_size: Maximum concurrent sequences (default: 32) enable_prefix_caching: Enable hash-based prefix caching (default: False) + enable_deepseek_mla_paging: Enable default-off DeepSeek V2/V3 MLA paging. + max_resident_paged_speculative_sessions: Resident paged-session cap. + min_free_mla_blocks_after_admission: Free blocks retained after admission. offload_dir: Path to offload directory (required) speculative_draft: Optional DFlash checkpoint, speculator, or draft module for greedy batch-1 serving. @@ -966,6 +1228,13 @@ def serve( max_batch_size=max_batch_size, enable_prefix_caching=enable_prefix_caching, speculative_draft=serving_speculator, + enable_deepseek_mla_paging=enable_deepseek_mla_paging, + max_resident_paged_speculative_sessions=( + max_resident_paged_speculative_sessions + ), + min_free_mla_blocks_after_admission=( + min_free_mla_blocks_after_admission + ), ) uvicorn = importlib.import_module("uvicorn") diff --git a/moe_infinity/entrypoints/openai/api_server_v2.py b/moe_infinity/entrypoints/openai/api_server_v2.py index 08cc9bf3..ba8b555b 100644 --- a/moe_infinity/entrypoints/openai/api_server_v2.py +++ b/moe_infinity/entrypoints/openai/api_server_v2.py @@ -483,6 +483,9 @@ def initialize_with_model( max_batch_size: int = 32, enable_prefix_caching: bool = False, speculative_draft: Optional[Any] = None, + enable_deepseek_mla_paging: bool = False, + max_resident_paged_speculative_sessions: int = 1, + min_free_mla_blocks_after_admission: int = 1, ) -> None: """Initialize the v2 server with a pre-loaded MoE model. @@ -502,6 +505,13 @@ def initialize_with_model( kv_cache_ratio=kv_cache_ratio, max_batch_size=max_batch_size, enable_prefix_caching=enable_prefix_caching, + enable_deepseek_mla_paging=enable_deepseek_mla_paging, + max_resident_paged_speculative_sessions=( + max_resident_paged_speculative_sessions + ), + min_free_mla_blocks_after_admission=( + min_free_mla_blocks_after_admission + ), ) engine_config = _build_engine_config(args=args, model=hf_model) @@ -1060,9 +1070,25 @@ async def _initialize_model() -> None: trust_remote_code=True, ) + enable_deepseek_mla_paging = bool( + getattr(args, "enable_deepseek_mla_paging", False) + ) + max_resident_paged_speculative_sessions = int( + getattr(args, "max_resident_paged_speculative_sessions", 1) + ) + min_free_mla_blocks_after_admission = int( + getattr(args, "min_free_mla_blocks_after_admission", 1) + ) moe_config = { "offload_path": os.path.join(args.offload_dir, args.model), "device_memory_ratio": args.device_memory_ratio, + "enable_deepseek_mla_paging": enable_deepseek_mla_paging, + "max_resident_paged_speculative_sessions": ( + max_resident_paged_speculative_sessions + ), + "min_free_mla_blocks_after_admission": ( + min_free_mla_blocks_after_admission + ), } if args.enable_prefix_caching: moe_config["enable_prefix_caching"] = True @@ -1857,6 +1883,15 @@ def _build_engine_config( "num_kv_heads": num_kv_heads, "head_dim": head_dim, "dtype": _resolve_dtype(model), + "enable_deepseek_mla_paging": bool( + getattr(args, "enable_deepseek_mla_paging", False) + ), + "max_resident_paged_speculative_sessions": int( + getattr(args, "max_resident_paged_speculative_sessions", 1) + ), + "min_free_mla_blocks_after_admission": int( + getattr(args, "min_free_mla_blocks_after_admission", 1) + ), } if eos_token_id is not None: config["eos_token_id"] = eos_token_id @@ -1882,6 +1917,28 @@ def parse_args() -> argparse.Namespace: parser.add_argument("--device-memory-ratio", type=float, default=0.75) parser.add_argument("--kv-cache-ratio", type=float, default=0.25) parser.add_argument("--max-batch-size", type=int, default=32) + parser.add_argument( + "--enable-deepseek-mla-paging", + action="store_true", + default=False, + help="Enable default-off DeepSeek V2/V3 paged MLA serving", + ) + parser.add_argument( + "--max-resident-paged-speculative-sessions", + type=int, + default=1, + help="Maximum concurrent resident paged-MLA speculative sessions", + ) + parser.add_argument( + "--min-free-mla-blocks-after-admission", + type=int, + default=1, + help=( + "Minimum MLA blocks that remain free after reserving all active " + "and newly admitted requests' full declared budgets plus maximum " + "transient DFlash verify peaks" + ), + ) parser.add_argument("--api-key", type=str, default=None) parser.add_argument("--rate-limit", type=int, default=0) parser.add_argument("--max-waiting-requests", type=int, default=0) diff --git a/moe_infinity/models/deepseek_mla_attention.py b/moe_infinity/models/deepseek_mla_attention.py new file mode 100644 index 00000000..77eec0a7 --- /dev/null +++ b/moe_infinity/models/deepseek_mla_attention.py @@ -0,0 +1,329 @@ +from __future__ import annotations + +# ruff: noqa: I001 + +from typing import Any, cast + +import torch +import torch.nn.functional as F + +try: + from transformers.models.deepseek_v2.modeling_deepseek_v2 import ( + DeepseekV2Attention, + apply_rotary_emb as apply_v2_rotary, + ) +except (ImportError, AttributeError): + DeepseekV2Attention = None # type: ignore[assignment,misc] + apply_v2_rotary = None + +try: + from transformers.models.deepseek_v3.modeling_deepseek_v3 import ( + DeepseekV3Attention, + apply_rotary_pos_emb, + apply_rotary_pos_emb_interleave, + ) +except (ImportError, AttributeError): + DeepseekV3Attention = None # type: ignore[assignment,misc] + apply_rotary_pos_emb = None + apply_rotary_pos_emb_interleave = None + +from moe_infinity.runtime.attention_types import AttentionMetadata +from moe_infinity.serving.mla_cache import MLAPagedKVCache + + +class _MLAPagedAttentionMixin: + _mla_cache: MLAPagedKVCache + _mla_metadata: AttentionMetadata | None = None + + def _mla_forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Any, + attention_mask: torch.Tensor | None, + *, + version: str, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + metadata = self._mla_metadata + if metadata is None: + raise RuntimeError("DeepSeek MLA paged context is not set") + if hidden_states.ndim != 3 or hidden_states.shape[0] != 1: + raise ValueError( + "DeepSeek MLA paged attention supports batch size 1" + ) + if self.layer_idx is None: + raise ValueError("DeepSeek MLA paged attention requires layer_idx") + + batch_size, query_len, _ = hidden_states.shape + if metadata.seq_lens.numel() != 1: + raise ValueError( + "DeepSeek MLA paged attention requires one seq_len" + ) + total_len = int(metadata.seq_lens.reshape(-1)[0].item()) + if total_len < query_len: + raise ValueError( + f"total_len {total_len} must be >= query_len {query_len}" + ) + if ( + metadata.slot_mapping.ndim != 1 + or metadata.slot_mapping.numel() != query_len + ): + raise ValueError( + f"slot_mapping must contain query_len {query_len} entries" + ) + if attention_mask is not None and attention_mask.shape[-1] < total_len: + raise ValueError( + f"attention_mask last dimension {attention_mask.shape[-1]} " + f"must be >= total_len {total_len}" + ) + block_table = metadata.block_tables.reshape(-1) + if metadata.seq_id is None: + raise ValueError( + "engine-owned DeepSeek MLA cache access requires metadata.seq_id" + ) + self._mla_cache.validate_owned_access( + metadata.seq_id, + block_table, + metadata.slot_mapping, + total_len, + ) + if self.q_lora_rank is None: + q_states = self.q_proj(hidden_states) + else: + q_states = self.q_b_proj( + self.q_a_layernorm(self.q_a_proj(hidden_states)) + ) + q_states = q_states.view( + batch_size, query_len, self.num_heads, self.qk_head_dim + ).transpose(1, 2) + q_nope, q_rope = torch.split( + q_states, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1 + ) + + compressed = self.kv_a_proj_with_mqa(hidden_states) + latent, k_rope = torch.split( + compressed, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1 + ) + latent = self.kv_a_layernorm(latent) + k_rope = k_rope.view(batch_size, 1, query_len, self.qk_rope_head_dim) + if version == "v2": + if apply_v2_rotary is None: + raise RuntimeError( + "Transformers DeepSeek V2 rotary helper unavailable" + ) + q_rope, k_rope = apply_v2_rotary( + q_rope, k_rope, position_embeddings.to(q_rope.device) + ) + else: + cos, sin = position_embeddings + if self.config.rope_interleave: + if apply_rotary_pos_emb_interleave is None: + raise RuntimeError( + "Transformers DeepSeek V3 interleaved RoPE unavailable" + ) + q_rope, k_rope = apply_rotary_pos_emb_interleave( + q_rope, k_rope, cos, sin + ) + else: + if apply_rotary_pos_emb is None: + raise RuntimeError( + "Transformers DeepSeek V3 RoPE unavailable" + ) + q_rope, k_rope = apply_rotary_pos_emb(q_rope, k_rope, cos, sin) + + self._mla_cache.write( + int(self.layer_idx), + latent.reshape(query_len, self.kv_lora_rank), + k_rope.reshape(query_len, self.qk_rope_head_dim), + metadata.slot_mapping, + ) + cached_latent, cached_rope = self._mla_cache.read( + int(self.layer_idx), block_table, total_len + ) + + expanded = ( + self.kv_b_proj(cached_latent) + .view( + 1, + total_len, + self.num_heads, + self.qk_nope_head_dim + self.v_head_dim, + ) + .transpose(1, 2) + ) + k_nope, values = torch.split( + expanded, [self.qk_nope_head_dim, self.v_head_dim], dim=-1 + ) + keys = torch.cat( + ( + k_nope, + cached_rope.view(1, 1, total_len, self.qk_rope_head_dim).expand( + 1, self.num_heads, total_len, self.qk_rope_head_dim + ), + ), + dim=-1, + ) + queries = torch.cat((q_nope, q_rope), dim=-1) + scores = torch.matmul(queries, keys.transpose(2, 3)) * float( + self.scaling + ) + past_len = total_len - query_len + causal = torch.arange(total_len, device=scores.device).view(1, 1, 1, -1) + limits = past_len + torch.arange(query_len, device=scores.device).view( + 1, 1, -1, 1 + ) + scores = scores.masked_fill( + causal > limits, torch.finfo(scores.dtype).min + ) + if attention_mask is not None: + scores = scores + attention_mask[..., :total_len] + weights = F.softmax(scores, dim=-1, dtype=torch.float32).to( + queries.dtype + ) + output = torch.matmul(weights, values).transpose(1, 2).contiguous() + output = self.o_proj(output.reshape(batch_size, query_len, -1)) + return output, weights + + +if DeepseekV2Attention is not None: + + class DeepseekV2MLAPagedAttention( + _MLAPagedAttentionMixin, DeepseekV2Attention + ): + def forward( + self, + hidden_states: torch.Tensor, + attention_mask: torch.Tensor | None = None, + past_key_values: Any = None, + position_embeddings: Any = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + del past_key_values, kwargs + return self._mla_forward( + hidden_states, position_embeddings, attention_mask, version="v2" + ) + +else: + DeepseekV2MLAPagedAttention = None # type: ignore[misc,assignment] + + +if DeepseekV3Attention is not None: + + class DeepseekV3MLAPagedAttention( + _MLAPagedAttentionMixin, DeepseekV3Attention + ): + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: Any, + attention_mask: torch.Tensor | None, + past_key_values: Any = None, + **kwargs: Any, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + del past_key_values, kwargs + return self._mla_forward( + hidden_states, position_embeddings, attention_mask, version="v3" + ) + +else: + DeepseekV3MLAPagedAttention = None # type: ignore[misc,assignment] + + +def adapt_deepseek_attention( + module: torch.nn.Module, + cache: MLAPagedKVCache, + *, + enabled: bool = False, +) -> torch.nn.Module: + """Adapt a real upstream attention object in place, preserving parameters.""" + if not enabled: + return module + target: type[torch.nn.Module] | None = None + if DeepseekV2Attention is not None and isinstance( + module, DeepseekV2Attention + ): + target = cast(type[torch.nn.Module], DeepseekV2MLAPagedAttention) + elif DeepseekV3Attention is not None and isinstance( + module, DeepseekV3Attention + ): + target = cast(type[torch.nn.Module], DeepseekV3MLAPagedAttention) + if target is None: + raise TypeError( + "module is not an installed upstream DeepSeek V2/V3 attention" + ) + if getattr(module, "layer_idx", None) is None: + raise ValueError("DeepSeek MLA adapter requires layer_idx") + module.__class__ = target + setattr(module, "_mla_cache", cache) + setattr(module, "_mla_metadata", None) + return module + + +def is_deepseek_mla_eligible(config: object, *, enabled: bool = False) -> bool: + """Fail-closed structural gate for the batch-one token-KV-only path.""" + if not enabled: + return False + model_type = str(getattr(config, "model_type", "")).lower() + if model_type not in {"deepseek_v2", "deepseek_v3"}: + return False + if not isinstance(getattr(config, "kv_lora_rank", None), int): + return False + for name in ( + "sliding_window", + "sliding_window_pattern", + "recurrent_chunk_size", + "hybrid_attention", + ): + value = getattr(config, name, None) + if value not in (None, False, 0): + return False + return True + + +def adapt_deepseek_model( + model: torch.nn.Module, + cache: MLAPagedKVCache, + *, + enabled: bool = False, +) -> list[torch.nn.Module]: + config = getattr(model, "config", None) + if config is None or not is_deepseek_mla_eligible(config, enabled=enabled): + return [] + adapted: list[torch.nn.Module] = [] + for module in model.modules(): + if DeepseekV2Attention is not None and isinstance( + module, DeepseekV2Attention + ): + adapted.append( + adapt_deepseek_attention(module, cache, enabled=True) + ) + elif DeepseekV3Attention is not None and isinstance( + module, DeepseekV3Attention + ): + adapted.append( + adapt_deepseek_attention(module, cache, enabled=True) + ) + return adapted + + +def set_deepseek_mla_context( + module: torch.nn.Module, metadata: AttentionMetadata +) -> None: + if not isinstance(module, _MLAPagedAttentionMixin): + raise TypeError("module is not an MLA paged attention adapter") + module._mla_metadata = metadata + + +def clear_deepseek_mla_context(module: torch.nn.Module) -> None: + if isinstance(module, _MLAPagedAttentionMixin): + module._mla_metadata = None + + +__all__ = [ + "DeepseekV2MLAPagedAttention", + "DeepseekV3MLAPagedAttention", + "adapt_deepseek_attention", + "adapt_deepseek_model", + "is_deepseek_mla_eligible", + "set_deepseek_mla_context", + "clear_deepseek_mla_context", +] diff --git a/moe_infinity/runtime/attention_types.py b/moe_infinity/runtime/attention_types.py index 35ba0410..445717c1 100644 --- a/moe_infinity/runtime/attention_types.py +++ b/moe_infinity/runtime/attention_types.py @@ -31,3 +31,4 @@ class AttentionMetadata: num_decode_tokens: int slot_mapping: torch.Tensor is_prefill: bool + seq_id: int | None = None diff --git a/moe_infinity/serving/__init__.py b/moe_infinity/serving/__init__.py index 75d47e72..a4ac27df 100644 --- a/moe_infinity/serving/__init__.py +++ b/moe_infinity/serving/__init__.py @@ -3,9 +3,16 @@ from .batch import BatchBuilder, BatchMetadata, SchedulerOutput from .engine import ContinuousBatchingEngine, RequestOutput from .kv_cache import BlockAllocator, BlockTable, PagedKVCache +from .mla_cache import MLAPagedKVCache from .sampler import Sampler from .scheduler import Scheduler from .sequence import SamplingParams, SequenceData, SequenceStatus +from .spec_session_driver import ( + EXECUTION_CONTEXT_TEMPORARY_DYNAMIC, + ServingSpecSession, + SpecSessionDriver, + TemporaryDynamicCacheContext, +) from .stream import StreamChunk, StreamManager RequestScheduler = Scheduler @@ -18,6 +25,7 @@ "BlockTable", "ContinuousBatchingEngine", "PagedKVCache", + "MLAPagedKVCache", "RequestOutput", "SamplingParams", "Sampler", @@ -26,6 +34,10 @@ "Sequence", "SequenceData", "SequenceStatus", + "EXECUTION_CONTEXT_TEMPORARY_DYNAMIC", + "ServingSpecSession", + "SpecSessionDriver", + "TemporaryDynamicCacheContext", "SchedulerOutput", "StreamChunk", "StreamManager", diff --git a/moe_infinity/serving/engine.py b/moe_infinity/serving/engine.py index b3f98663..1beb271a 100644 --- a/moe_infinity/serving/engine.py +++ b/moe_infinity/serving/engine.py @@ -1,5 +1,6 @@ from __future__ import annotations +import logging from collections.abc import Iterable, Iterator from dataclasses import dataclass from math import ceil @@ -7,7 +8,12 @@ import torch -from .batch import BatchBuilder, BatchMetadata, split_prefill_decode_batch +from .batch import ( + BatchBuilder, + BatchMetadata, + _slice_batch, + split_prefill_decode_batch, +) from .kv_cache import PagedKVCache from .memory_manager import MemoryManager from .model_runner import ModelRunner @@ -19,6 +25,12 @@ SequenceGroup, SequenceStatus, ) +from .spec_session_driver import ( + ServingSpecSession, + SpecSessionDriver, +) + +logger = logging.getLogger(__name__) class EvictionSyncAdapter(Protocol): @@ -44,6 +56,29 @@ def generate( _VERIFY_ADMISSION_MAX_RETRIES = 1024 +def _debug_cleanup_reporting_failure( + *, + action: str, + request_id: str, + phase: str, + primary: BaseException, + reporting_error: BaseException, +) -> None: + try: + logger.debug( + "cleanup %s attachment failed for request %s during %s; " + "primary=%s reporting=%s", + action, + request_id, + phase, + type(primary).__name__, + type(reporting_error).__name__, + exc_info=reporting_error, + ) + except BaseException: + return + + def set_eviction_sync(adapter: Optional[EvictionSyncAdapter]) -> None: global _eviction_sync _eviction_sync = adapter @@ -151,9 +186,38 @@ def __init__( self.sampler = Sampler() self.batch_builder = BatchBuilder() self.speculative_draft = speculative_draft + enable_paged_mla = self.config.get("enable_deepseek_mla_paging", False) + if not isinstance(enable_paged_mla, bool): + raise ValueError( + "enable_deepseek_mla_paging must be a boolean value" + ) + max_resident_paged_sessions = self._get_int_config( + "max_resident_paged_speculative_sessions", 1 + ) + min_free_mla_blocks = self._get_int_config( + "min_free_mla_blocks_after_admission", 1 + ) + if max_resident_paged_sessions < 0: + raise ValueError( + "max_resident_paged_speculative_sessions must be >= 0" + ) + if min_free_mla_blocks < 1: + raise ValueError("min_free_mla_blocks_after_admission must be >= 1") self._verify_scheduling_enabled = ( self.scheduler.verify_scheduling_enabled ) + self._spec_session_driver = ( + SpecSessionDriver( + speculative_draft, + enable_paged_mla=enable_paged_mla, + max_resident_paged_speculative_sessions=( + max_resident_paged_sessions + ), + min_free_mla_blocks_after_admission=min_free_mla_blocks, + ) + if self._can_drive_verify_rounds() + else None + ) self._next_seq_id = 0 self._sequences: dict[int, SequenceData] = {} @@ -163,6 +227,7 @@ def __init__( self._callbacks: dict[str, list[Callable[[RequestOutput], None]]] = {} self._completed_request_ids: set[str] = set() self._cancelled_request_ids: set[str] = set() + self._request_failures: dict[str, dict[str, str]] = {} self._num_steps = 0 self._total_generated_tokens = 0 @@ -206,12 +271,18 @@ def add_request( self.scheduler.add_request(group) def step(self) -> list[RequestOutput]: + self._prepare_speculative_rounds() scheduler_output = self.scheduler.schedule() + outputs = self._verify_speculative_rounds( + scheduler_output.verify_seq_ids + ) if ( not scheduler_output.prefill_seq_ids and not scheduler_output.decode_seq_ids ): - return [] + if outputs: + self._num_steps += 1 + return outputs batch = self.batch_builder.from_scheduler_output( scheduler_output, @@ -223,11 +294,58 @@ def step(self) -> list[RequestOutput]: "scheduler produced an empty batch; empty prompts are not supported" ) - if self._can_delegate_speculative(batch): + if self._spec_session_driver is None and self._can_delegate_speculative( + batch + ): if self._can_drive_verify_rounds(): return self._step_speculative_session(batch) return self._step_speculative(batch) + # Compatibility limit: pre-Stage4a session doubles do not accept a + # request-scoped generator and retain the original singleton, + # whole-request Step-5 behavior. Canonical SpecSession implementations + # always take the persistent path below. + if ( + self._spec_session_driver is not None + and not self._spec_session_driver.supports_request_generator + and self._can_delegate_speculative(batch) + ): + outputs.extend(self._step_speculative_session(batch)) + return outputs + + speculative_indices = [ + index + for index, seq_id in enumerate(batch.seq_ids) + if self._can_start_persistent_speculative( + seq_id, batch.is_prefill[index] + ) + ] + for index in speculative_indices: + outputs.extend( + self._begin_persistent_speculative(batch.seq_ids[index]) + ) + + speculative_index_set = set(speculative_indices) + fallback_indices = [ + index + for index in range(len(batch.seq_ids)) + if index not in speculative_index_set + ] + if fallback_indices: + fallback_batch = ( + batch + if len(fallback_indices) == len(batch.seq_ids) + else _slice_batch(batch, fallback_indices) + ) + outputs.extend(self._step_standard(fallback_batch)) + + if speculative_indices and not fallback_indices: + self._num_steps += 1 + return outputs + + def _step_standard(self, batch: BatchMetadata) -> list[RequestOutput]: + """Execute the ordinary serving path for a scheduler-selected subset.""" + logits = self._execute_batch(batch) last_token_logits = self._extract_last_token_logits(logits, batch) sampler_output = self.sampler.sample( @@ -304,6 +422,296 @@ def step(self) -> list[RequestOutput]: return outputs + @property + def speculative_sessions(self) -> dict[int, ServingSpecSession]: + """Live Stage 4a records, keyed by serving sequence id.""" + driver = self._spec_session_driver + return {} if driver is None else driver.sessions + + def _can_start_persistent_speculative( + self, seq_id: int, is_prefill: bool + ) -> bool: + """Check semantics that the temporary DynamicCache path can preserve.""" + if self._spec_session_driver is None or not is_prefill: + return False + sequence = self._sequences[seq_id] + params = sequence.sampling_params + return ( + sequence.status is SequenceStatus.PREFILL + and not sequence.output_token_ids + and params.max_tokens > 0 + and not params.stop + and params.temperature >= 0 + and params.top_p > 0 + and params.top_p <= 1 + and params.repetition_penalty == 1.0 + and params.logprobs <= 0 + and not self._has_unsupported_speculative_metadata(params) + ) + + @staticmethod + def _has_unsupported_speculative_metadata(params: SamplingParams) -> bool: + for name in ( + "grammar", + "guided_decoding", + "response_format", + "logit_bias", + "logits_processors", + ): + if getattr(params, name, None): + return True + for name in ("presence_penalty", "frequency_penalty", "min_p"): + if float(getattr(params, name, 0.0) or 0.0) != 0.0: + return True + return False + + def _begin_persistent_speculative(self, seq_id: int) -> list[RequestOutput]: + driver = self._spec_session_driver + if driver is None: + raise RuntimeError( + "persistent speculative driver is not configured" + ) + sequence = self._sequences[seq_id] + request_id = self._sequence_to_request_id[seq_id] + params = sequence.sampling_params + stop_token_ids = ( + [self.eos_token_id] if self.eos_token_id is not None else [] + ) + generator: torch.Generator | None = None + if params.temperature > 0: + generator_device = ( + self.device + if self.device.type == "cuda" + else torch.device("cpu") + ) + generator = torch.Generator(device=generator_device) + base_seed = self._get_int_config("speculative_seed", 0) + generator.manual_seed(base_seed + seq_id) + record = driver.begin( + request_id=request_id, + seq_id=seq_id, + prompt_token_ids=sequence.prompt_token_ids, + max_new_tokens=params.max_tokens, + temperature=params.temperature, + top_k=max(0, params.top_k), + top_p=params.top_p, + stop_token_ids=stop_token_ids, + callbacks=tuple(self._callbacks.get(request_id, ())), + generator=generator, + ) + sequence.set_status(SequenceStatus.DRAFT) + committed = driver.commit(record) + return self._publish_speculative_commit(record, committed) + + def _prepare_speculative_rounds(self) -> None: + driver = self._spec_session_driver + if driver is None: + return + for record in tuple(driver.sessions.values()): + sequence = self._sequences.get(record.seq_id) + if ( + record.cancelled + or record.released + or sequence is None + or sequence.status is not SequenceStatus.DRAFT + or record.finished + ): + continue + try: + draft = driver.draft(record) + except BaseException as exc: + self._fail_speculative_request(record, "draft", exc) + raise + if record.cancelled or record.released: + continue + self.scheduler.set_verify_demand( + record.seq_id, + tokens=int(getattr(draft, "tokens")), + expert_bytes=int(getattr(draft, "expert_bytes")), + in_flight=False, + ) + + def _verify_speculative_rounds( + self, admitted_seq_ids: list[int] + ) -> list[RequestOutput]: + driver = self._spec_session_driver + if driver is None: + return [] + outputs: list[RequestOutput] = [] + for seq_id in admitted_seq_ids: + record = driver.sessions.get(seq_id) + sequence = self._sequences.get(seq_id) + if ( + record is None + or sequence is None + or record.pending_draft is None + ): + continue + sequence.set_status(SequenceStatus.VERIFY) + try: + _ = driver.verify(record) + except BaseException as exc: + self._fail_speculative_request(record, "verify", exc) + raise + self.scheduler.clear_verify_demand(seq_id) + if ( + record.cancelled + or record.released + or seq_id not in self._sequences + ): + continue + committed = driver.commit(record) + outputs.extend(self._publish_speculative_commit(record, committed)) + return outputs + + def _fail_speculative_request( + self, + failed_record: ServingSpecSession, + phase: str, + primary: BaseException, + ) -> None: + """Fail one request atomically while preserving its backend exception.""" + request_id = failed_record.request_id + failure = { + "phase": phase, + "failure_type": type(primary).__name__, + "code": f"speculative_{phase}_failed", + } + self._request_failures[request_id] = failure + seq_ids = list(self._request_to_seq_ids.get(request_id, ())) + cleanup_errors: list[BaseException] = [] + driver = self._spec_session_driver + if driver is not None: + records = [ + record + for record in tuple(driver.sessions.values()) + if record.request_id == request_id + ] + for record in records: + self.scheduler.clear_verify_demand(record.seq_id) + try: + driver.fail(record, failure) + except BaseException as exc: + cleanup_errors.append(exc) + try: + self.scheduler.abort_request(request_id) + except BaseException as exc: + cleanup_errors.append(exc) + + self._callbacks.pop(request_id, None) + self._request_outputs.pop(request_id, None) + self._request_to_seq_ids.pop(request_id, None) + for seq_id in seq_ids: + self.scheduler.clear_verify_demand(seq_id) + self._sequence_to_request_id.pop(seq_id, None) + self._sequences.pop(seq_id, None) + if _eviction_sync is not None: + try: + _eviction_sync.on_request_aborted(request_id) + except BaseException as exc: + cleanup_errors.append(exc) + if cleanup_errors: + cleanup_metadata = tuple(cleanup_errors) + try: + setattr(primary, "session_cleanup_errors", cleanup_metadata) + except BaseException as reporting_error: + _debug_cleanup_reporting_failure( + action="metadata", + request_id=request_id, + phase=phase, + primary=primary, + reporting_error=reporting_error, + ) + try: + add_note = getattr(primary, "add_note", None) + except BaseException as reporting_error: + _debug_cleanup_reporting_failure( + action="note-lookup", + request_id=request_id, + phase=phase, + primary=primary, + reporting_error=reporting_error, + ) + add_note = None + if callable(add_note): + for cleanup_error in cleanup_metadata: + try: + add_note( + "speculative request cleanup failed: " + f"{cleanup_error}" + ) + except BaseException as reporting_error: + _debug_cleanup_reporting_failure( + action="note", + request_id=request_id, + phase=phase, + primary=primary, + reporting_error=reporting_error, + ) + + def _publish_speculative_commit( + self, + record: ServingSpecSession, + committed: tuple[int, ...], + ) -> list[RequestOutput]: + if record.cancelled or record.released: + return [] + sequence = self._sequences.get(record.seq_id) + if sequence is None: + return [] + + outputs: list[RequestOutput] = [] + for token_id in committed: + sequence.append_output_token(token_id) + self._request_outputs[record.request_id][record.seq_id].append( + token_id + ) + self._total_generated_tokens += 1 + finish_reason = self._get_finish_reason(sequence, token_id) + finished = finish_reason is not None + output = RequestOutput( + request_id=record.request_id, + seq_id=record.seq_id, + token_id=token_id, + token_text=self._decode_token(token_id), + finished=finished, + finish_reason=finish_reason, + usage=(self._build_usage(sequence) if finished else None), + ) + outputs.append(output) + for callback in record.callbacks: + callback(output) + if finished: + break + + finished = record.finished or self._output_finished(outputs) + committed_count = len(outputs) + if finished: + if sequence.status in (SequenceStatus.DRAFT, SequenceStatus.VERIFY): + sequence.set_status(SequenceStatus.FINISHED) + self.scheduler.update_after_step( + completed_seq_ids=[record.seq_id], + new_decode_seq_ids=[], + committed_counts={record.seq_id: committed_count}, + ) + driver = self._spec_session_driver + if driver is not None: + driver.release(record) + if self._is_request_finished(record.request_id): + self._completed_request_ids.add(record.request_id) + if _eviction_sync is not None: + _eviction_sync.on_request_finished(record.request_id) + self._callbacks.pop(record.request_id, None) + else: + if sequence.status is SequenceStatus.VERIFY: + sequence.set_status(SequenceStatus.DRAFT) + self.scheduler.update_after_step( + completed_seq_ids=[], + new_decode_seq_ids=[], + committed_counts={record.seq_id: committed_count}, + ) + return outputs + def _can_delegate_speculative(self, batch: BatchMetadata) -> bool: """Whether this fresh singleton request can use the proven sync loop. @@ -624,12 +1032,19 @@ def abort_request(self, request_id: str) -> None: SequenceStatus.WAITING, SequenceStatus.PREFILL, SequenceStatus.DECODE, + SequenceStatus.DRAFT, + SequenceStatus.VERIFY, SequenceStatus.SWAPPED, } for seq_id in seq_ids if seq_id in self._sequences ) + driver = self._spec_session_driver + if driver is not None: + for seq_id in seq_ids: + driver.cancel(seq_id) + self.scheduler.clear_verify_demand(seq_id) self.scheduler.abort_request(request_id) _ = self._callbacks.pop(request_id, None) @@ -659,14 +1074,34 @@ def get_stats(self) -> dict[str, object]: "pending_requests": len(self._pending_request_ids()), "completed_requests": len(self._completed_request_ids), "cancelled_requests": len(self._cancelled_request_ids), + "failed_requests": len(self._request_failures), "num_steps": self._num_steps, "total_generated_tokens": self._total_generated_tokens, "kv_cache_num_blocks": self.kv_cache.num_blocks, "kv_cache_free_blocks": self.kv_cache.block_allocator.num_free_blocks, "sequence_status_counts": status_counts, + "speculative_execution_context": ( + self._spec_session_driver.execution_context_mode + if self._spec_session_driver is not None + else None + ), + "speculative_sessions": [ + record.diagnostics() + for record in self.speculative_sessions.values() + ], + "paged_mla_admission": ( + self._spec_session_driver.admission_stats + if self._spec_session_driver is not None + else None + ), "memory": self.memory_manager.report(), } + def get_request_failure(self, request_id: str) -> dict[str, str]: + if request_id not in self._request_failures: + raise KeyError(f"request_id '{request_id}' has no recorded failure") + return dict(self._request_failures[request_id]) + def get_config(self) -> dict[str, object]: config: dict[str, object] = {} for key, value in self.config.items(): @@ -808,6 +1243,8 @@ def _pending_request_ids(self) -> list[str]: SequenceStatus.WAITING, SequenceStatus.PREFILL, SequenceStatus.DECODE, + SequenceStatus.DRAFT, + SequenceStatus.VERIFY, SequenceStatus.SWAPPED, } diff --git a/moe_infinity/serving/mla_cache.py b/moe_infinity/serving/mla_cache.py new file mode 100644 index 00000000..04556d63 --- /dev/null +++ b/moe_infinity/serving/mla_cache.py @@ -0,0 +1,226 @@ +from __future__ import annotations + +from dataclasses import dataclass, field + +import torch + +from .kv_cache import BlockAllocator, BlockTable + + +@dataclass +class MLAPagedKVCache: + """Engine-owned DeepSeek MLA pages stored as ``[kv_c_normed | k_pe]``.""" + + num_blocks: int + block_size: int + num_layers: int + latent_dim: int + rope_dim: int + dtype: torch.dtype + device: torch.device | None = None + block_allocator: BlockAllocator = field(init=False) + _sequence_tables: dict[int, BlockTable] = field( + init=False, default_factory=dict + ) + _mla_cache: torch.Tensor = field(init=False) + + def __post_init__(self) -> None: + for name in ( + "num_blocks", + "block_size", + "num_layers", + "latent_dim", + "rope_dim", + ): + if int(getattr(self, name)) <= 0: + raise ValueError(f"{name} must be > 0") + if self.device is None: + self.device = torch.device( + "cuda" if torch.cuda.is_available() else "cpu" + ) + elif self.device.type == "cuda" and not torch.cuda.is_available(): + self.device = torch.device("cpu") + self.block_allocator = BlockAllocator( + self.num_blocks, self.block_size, self.device + ) + self._mla_cache = torch.zeros( + self.num_layers, + self.num_blocks, + self.block_size, + self.latent_dim + self.rope_dim, + dtype=self.dtype, + device=self.device, + ) + + def allocate_sequence(self, seq_id: int, num_tokens: int) -> None: + if seq_id in self._sequence_tables: + raise ValueError(f"sequence {seq_id} already exists") + if num_tokens < 0: + raise ValueError("num_tokens must be >= 0") + table = BlockTable(self.block_allocator) + for _ in range(num_tokens): + table.append_token() + self._sequence_tables[seq_id] = table + + def append_tokens(self, seq_id: int, num_new_tokens: int) -> None: + if num_new_tokens < 0: + raise ValueError("num_new_tokens must be >= 0") + table = self._require_sequence(seq_id) + for _ in range(num_new_tokens): + table.append_token() + + def truncate_tokens(self, seq_id: int, new_len: int) -> None: + table = self._require_sequence(seq_id) + current = table.num_computed_tokens() + if new_len < 0 or new_len > current: + raise ValueError( + "truncate_tokens requires 0 <= new_len <= current length" + ) + blocks_needed = (new_len + self.block_size - 1) // self.block_size + block_ids = table.get_block_ids() + freed = block_ids[blocks_needed:] + if freed: + self.block_allocator.free(freed) + table.restore_blocks(block_ids[:blocks_needed], new_len) + + def free_sequence(self, seq_id: int) -> None: + table = self._sequence_tables.pop(seq_id, None) + if table is not None: + table.release() + + def get_block_table(self, seq_id: int) -> list[int]: + return self._require_sequence(seq_id).get_block_ids() + + def get_mla_cache_tensors(self) -> torch.Tensor: + return self._mla_cache + + @property + def free_block_count(self) -> int: + """Return currently unowned blocks without exposing allocator mutation.""" + return int(self.block_allocator.num_free_blocks) + + def validate_owned_access( + self, + seq_id: int, + block_table: torch.Tensor, + slot_mapping: torch.Tensor, + total_len: int, + ) -> None: + table = self._require_sequence(seq_id) + if total_len > table.num_computed_tokens(): + raise ValueError( + f"total_len {total_len} exceeds allocated sequence length " + f"{table.num_computed_tokens()} for seq_id {seq_id}" + ) + needed = (total_len + self.block_size - 1) // self.block_size + expected = torch.tensor( + table.get_block_ids()[:needed], + dtype=torch.long, + device=block_table.device, + ) + actual = block_table.reshape(-1)[:needed].to(dtype=torch.long) + if actual.numel() != expected.numel() or not torch.equal( + actual, expected + ): + raise ValueError( + f"block_tables do not match allocated sequence {seq_id} pages" + ) + slot_pages = torch.div( + slot_mapping.to(device=expected.device, dtype=torch.long), + self.block_size, + rounding_mode="floor", + ) + if slot_pages.numel() and not bool( + torch.isin(slot_pages, expected).all().item() + ): + raise ValueError( + f"slot_mapping references pages not owned by sequence {seq_id}" + ) + + def write( + self, + layer_idx: int, + latent: torch.Tensor, + rope: torch.Tensor, + slot_mapping: torch.Tensor, + ) -> None: + self._validate_layer(layer_idx) + if latent.ndim != 2 or latent.shape[1] != self.latent_dim: + raise ValueError("latent must have shape [num_tokens, latent_dim]") + if rope.ndim != 2 or rope.shape != (latent.shape[0], self.rope_dim): + raise ValueError("rope must have shape [num_tokens, rope_dim]") + if slot_mapping.ndim != 1 or slot_mapping.shape[0] != latent.shape[0]: + raise ValueError("slot_mapping must have shape [num_tokens]") + packed = torch.cat((latent, rope), dim=-1).to( + device=self.device, dtype=self.dtype + ) + slots = slot_mapping.to(device=self.device, dtype=torch.long) + if slots.numel() == 0: + return + if bool((slots < 0).any().item()): + raise ValueError("slot_mapping contains negative slot") + if bool((slots >= self.num_blocks * self.block_size).any().item()): + raise ValueError("slot_mapping points past allocated pages") + + # Advanced assignment with duplicate indices has backend-dependent write + # ordering. Select each slot's final source row first so semantics match + # the former left-to-right loop deterministically on CPU and GPU. + unique_slots, inverse = torch.unique( + slots, sorted=False, return_inverse=True + ) + source_rows = torch.arange(slots.numel(), device=self.device) + final_rows = torch.full_like(unique_slots, -1) + final_rows.scatter_reduce_( + 0, inverse, source_rows, reduce="amax", include_self=True + ) + final_slots = slots[final_rows] + pages = torch.div(final_slots, self.block_size, rounding_mode="floor") + offsets = torch.remainder(final_slots, self.block_size) + self._mla_cache[layer_idx, pages, offsets] = packed[final_rows] + + def read( + self, + layer_idx: int, + block_table: list[int] | torch.Tensor, + seq_len: int, + ) -> tuple[torch.Tensor, torch.Tensor]: + self._validate_layer(layer_idx) + if seq_len < 0: + raise ValueError("seq_len must be >= 0") + page_ids = torch.as_tensor( + block_table, device=self.device, dtype=torch.long + ).reshape(-1) + needed = (seq_len + self.block_size - 1) // self.block_size + if page_ids.numel() < needed: + raise ValueError("block_table is too short for seq_len") + selected_pages = page_ids[:needed] + if selected_pages.numel() and bool( + ((selected_pages < 0) | (selected_pages >= self.num_blocks)) + .any() + .item() + ): + raise ValueError( + "block_table contains page outside allocated cache" + ) + if seq_len == 0: + packed = self._mla_cache.new_empty( + (0, self.latent_dim + self.rope_dim) + ) + else: + packed = self._mla_cache[layer_idx, selected_pages].reshape( + -1, self.latent_dim + self.rope_dim + )[:seq_len] + return packed[:, : self.latent_dim], packed[:, self.latent_dim :] + + def _validate_layer(self, layer_idx: int) -> None: + if not 0 <= int(layer_idx) < self.num_layers: + raise IndexError(f"layer_idx {layer_idx} is outside the MLA cache") + + def _require_sequence(self, seq_id: int) -> BlockTable: + try: + return self._sequence_tables[seq_id] + except KeyError as error: + raise KeyError(f"unknown sequence id: {seq_id}") from error + + +__all__ = ["MLAPagedKVCache"] diff --git a/moe_infinity/serving/model_runner.py b/moe_infinity/serving/model_runner.py index 7fb366cc..f4dc69d5 100644 --- a/moe_infinity/serving/model_runner.py +++ b/moe_infinity/serving/model_runner.py @@ -94,11 +94,43 @@ def execute( batch: BatchMetadata, past_key_values: object = None, ) -> torch.Tensor: + return self._execute(batch, past_key_values=past_key_values, rich=False) + + def execute_rich( + self, + batch: BatchMetadata, + *, + cache_handles: tuple[object, ...] = (), + ) -> object: + """Run the normal serving forward and retain speculative state.""" + return self._execute(batch, rich=True, cache_handles=cache_handles) + + def _execute( + self, + batch: BatchMetadata, + past_key_values: object = None, + *, + rich: bool, + cache_handles: tuple[object, ...] = (), + ) -> Any: self._configure_expert_tracing(len(batch.seq_ids)) self._advance_request_id() if batch.total_tokens == 0: - return self._empty_logits() + logits = self._empty_logits() + if not rich: + return logits + from moe_infinity.spec_decode.protocols import RichForwardResult + + shared = cache_handles[0] if cache_handles else None + return RichForwardResult( + logits=logits, + hidden_states=(), + cache_handle=shared, + cache_handles=cache_handles or (shared,) * len(batch.seq_ids), + row_offsets=tuple(batch.token_offsets), + row_lengths=tuple(batch.seq_lengths), + ) model_inputs = self.prepare_inputs(batch) @@ -110,6 +142,8 @@ def execute( **model_inputs, "use_cache": True, } + if rich: + forward_kwargs["output_hidden_states"] = True if past_key_values is not None: forward_kwargs["past_key_values"] = past_key_values @@ -151,7 +185,31 @@ def execute( raise ValueError( f"packed logits row count must match batch.total_tokens; got {logits.size(0)} vs {batch.total_tokens}" ) - return logits + if not rich: + return logits + + from moe_infinity.spec_decode.protocols import RichForwardResult + + hidden_value = getattr(outputs, "hidden_states", None) + if hidden_value is None: + raise ValueError("rich model output must include hidden_states") + token_mask = model_inputs["attention_mask"].to(dtype=torch.bool) + hidden_states = tuple( + hidden[token_mask.to(device=hidden.device)] + if hidden.dim() == 3 + else hidden + for hidden in hidden_value + ) + cache_handle = getattr(outputs, "past_key_values", None) + row_handles = cache_handles or (cache_handle,) * len(batch.seq_ids) + return RichForwardResult( + logits=logits, + hidden_states=hidden_states, + cache_handle=cache_handle, + cache_handles=row_handles, + row_offsets=tuple(batch.token_offsets), + row_lengths=tuple(batch.seq_lengths), + ) def _configure_expert_tracing(self, num_sequences: int) -> None: tracer = getattr(self.engine, "expert_tracer", None) diff --git a/moe_infinity/serving/scheduler.py b/moe_infinity/serving/scheduler.py index d29e82db..b021dae8 100644 --- a/moe_infinity/serving/scheduler.py +++ b/moe_infinity/serving/scheduler.py @@ -414,6 +414,8 @@ def update_after_step( SequenceStatus.WAITING, SequenceStatus.PREFILL, SequenceStatus.DECODE, + SequenceStatus.DRAFT, + SequenceStatus.VERIFY, SequenceStatus.SWAPPED, ): sequence.set_status(SequenceStatus.FINISHED) @@ -467,22 +469,32 @@ def get_running_seq_ids(self) -> list[int]: if sequence.status in ( SequenceStatus.PREFILL, SequenceStatus.DECODE, + SequenceStatus.DRAFT, + SequenceStatus.VERIFY, ): running_seq_ids.append(sequence.seq_id) return running_seq_ids def _preempt_oldest_running_group(self) -> list[int]: + preserved: list[SequenceGroup] = [] while self._running: group = self._running.popleft() + active = [ + sequence + for sequence in group.sequences + if sequence.status + not in (SequenceStatus.FINISHED, SequenceStatus.CANCELLED) + ] + if not active or any( + sequence.status + not in (SequenceStatus.PREFILL, SequenceStatus.DECODE) + for sequence in active + ): + preserved.append(group) + continue preempted_seq_ids: list[int] = [] - for sequence in group.sequences: - if sequence.status not in ( - SequenceStatus.PREFILL, - SequenceStatus.DECODE, - ): - continue - + for sequence in active: try: self.kv_cache.swap_out(sequence.seq_id) except KeyError: @@ -493,8 +505,10 @@ def _preempt_oldest_running_group(self) -> list[int]: if preempted_seq_ids: self._swapped.append(group) + self._running.extendleft(reversed(preserved)) return preempted_seq_ids + self._running.extend(preserved) return [] def _recover_swapped_groups( diff --git a/moe_infinity/serving/spec_cache_adapter.py b/moe_infinity/serving/spec_cache_adapter.py new file mode 100644 index 00000000..7ac1852f --- /dev/null +++ b/moe_infinity/serving/spec_cache_adapter.py @@ -0,0 +1,136 @@ +from __future__ import annotations + +from dataclasses import dataclass + +import torch + +from moe_infinity.runtime.attention_types import AttentionMetadata +from moe_infinity.spec_decode.protocols import CacheSnapshot + +from .mla_cache import MLAPagedKVCache + +EXECUTION_CONTEXT_PAGED_MLA = "paged_mla" + + +@dataclass(frozen=True) +class PagedCacheSnapshot(CacheSnapshot): + block_table: tuple[int, ...] + + +class PagedCacheAdapter: + """Per-sequence speculative handle over one engine-owned MLA cache.""" + + cache_kind = "paged" + mode = EXECUTION_CONTEXT_PAGED_MLA + + def __init__( + self, + cache: MLAPagedKVCache, + seq_id: int, + initial_length: int, + ) -> None: + if initial_length < 0: + raise ValueError("initial_length must be >= 0") + self.cache = cache + self.seq_id = int(seq_id) + self._logical_length = int(initial_length) + self._released = False + cache.allocate_sequence(self.seq_id, self._logical_length) + + def _ensure_active(self) -> None: + if self._released: + raise RuntimeError("paged cache adapter has been released") + + def snapshot(self) -> PagedCacheSnapshot: + self._ensure_active() + return PagedCacheSnapshot( + logical_length=self._logical_length, + block_table=tuple(self.cache.get_block_table(self.seq_id)), + ) + + def restore(self, snapshot: CacheSnapshot) -> None: + self._ensure_active() + if not isinstance(snapshot, PagedCacheSnapshot): + raise TypeError("paged cache restore requires PagedCacheSnapshot") + current = tuple(self.cache.get_block_table(self.seq_id)) + if current[: len(snapshot.block_table)] != snapshot.block_table: + raise RuntimeError( + "paged cache block-table prefix changed since snapshot" + ) + self.truncate(snapshot.logical_length) + + def append(self, token_count: int) -> None: + self._ensure_active() + if token_count < 0: + raise ValueError("token_count must be >= 0") + self.cache.append_tokens(self.seq_id, token_count) + self._logical_length += token_count + + def truncate(self, logical_length: int) -> None: + self._ensure_active() + if logical_length < 0 or logical_length > self._logical_length: + raise ValueError( + "logical_length must be between 0 and the current logical length" + ) + self.cache.truncate_tokens(self.seq_id, logical_length) + self._logical_length = logical_length + + def logical_length(self) -> int: + self._ensure_active() + return self._logical_length + + def build_attention_metadata( + self, *, query_length: int, is_prefill: bool + ) -> AttentionMetadata: + self._ensure_active() + if query_length < 0 or query_length > self._logical_length: + raise ValueError( + "query_length must be between 0 and the allocated logical length" + ) + block_table = self.cache.get_block_table(self.seq_id) + device = self.cache.device + table_tensor = torch.tensor( + [block_table], dtype=torch.int32, device=device + ) + start = self._logical_length - query_length + slots = [ + block_table[position // self.cache.block_size] + * self.cache.block_size + + position % self.cache.block_size + for position in range(start, self._logical_length) + ] + return AttentionMetadata( + block_tables=table_tensor, + seq_lens=torch.tensor( + [self._logical_length], dtype=torch.int32, device=device + ), + max_seq_len=self._logical_length, + num_prefill_tokens=query_length if is_prefill else 0, + num_decode_tokens=0 if is_prefill else query_length, + slot_mapping=torch.tensor(slots, dtype=torch.int64, device=device), + is_prefill=is_prefill, + seq_id=self.seq_id, + ) + + def swap_out(self) -> bool: + """MLA pages currently stay resident during scheduler preemption.""" + self._ensure_active() + return False + + def swap_in(self) -> bool: + self._ensure_active() + return False + + def release(self) -> None: + if self._released: + return + self.cache.free_sequence(self.seq_id) + self._logical_length = 0 + self._released = True + + +__all__ = [ + "EXECUTION_CONTEXT_PAGED_MLA", + "PagedCacheAdapter", + "PagedCacheSnapshot", +] diff --git a/moe_infinity/serving/spec_session_driver.py b/moe_infinity/serving/spec_session_driver.py new file mode 100644 index 00000000..da05aee7 --- /dev/null +++ b/moe_infinity/serving/spec_session_driver.py @@ -0,0 +1,662 @@ +"""Temporary DynamicCache execution records for continuous speculative serving. + +Stage 4a keeps lifecycle and scheduling in ``ContinuousBatchingEngine`` while +each eligible sequence owns a persistent canonical ``SpecSession``. The cache +context here is intentionally private and temporary; it is not the serving +``PagedKVCache`` ownership contract planned for Stage 4b. +""" + +from __future__ import annotations + +from collections.abc import Callable, Iterator, Sequence +from contextlib import contextmanager +from dataclasses import dataclass, field +from inspect import Parameter, signature +from typing import TYPE_CHECKING, Any, cast + +import torch + +from .mla_cache import MLAPagedKVCache +from .spec_cache_adapter import ( + EXECUTION_CONTEXT_PAGED_MLA, + PagedCacheAdapter, +) +from .spec_state import SpecDecodeState + +if TYPE_CHECKING: + from moe_infinity.spec_decode.dflash import SpecSession + + +EXECUTION_CONTEXT_TEMPORARY_DYNAMIC = "temporary_dynamic" + + +@dataclass +class TemporaryDynamicCacheContext: + """Explicit Stage 4a bridge around one session's private dense caches.""" + + owner: object | None + target_cache: object | None + draft_cache: object | None + mode: str = EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + released: bool = False + owned_caches: list[object] = field(default_factory=list) + + def __post_init__(self) -> None: + self.refresh(self.target_cache, self.draft_cache) + + def refresh( + self, target_cache: object | None, draft_cache: object | None + ) -> None: + """Release superseded caches and retain only currently active objects.""" + active = [ + cache for cache in (target_cache, draft_cache) if cache is not None + ] + self.target_cache = target_cache + self.draft_cache = draft_cache + errors: list[BaseException] = [] + retained: list[object] = [] + for cache in self.owned_caches: + if any(cache is current for current in active): + retained.append(cache) + continue + crop = getattr(cache, "crop", None) + if callable(crop): + try: + crop(0) + except BaseException as exc: + errors.append(exc) + retained.append(cache) + self.owned_caches = retained + for cache in active: + if cache is not None and not any( + cache is owned for owned in self.owned_caches + ): + self.owned_caches.append(cache) + if errors: + raise errors[0] + + @contextmanager + def activate(self) -> Iterator[None]: + if self.released: + raise RuntimeError( + "temporary DynamicCache context has been released" + ) + previous: object | None = None + if self.owner is not None: + previous = getattr(self.owner, "_cached_past_key_values", None) + setattr(self.owner, "_cached_past_key_values", self.target_cache) + try: + yield + finally: + if self.owner is not None: + setattr(self.owner, "_cached_past_key_values", previous) + + def release(self) -> None: + if self.released: + return + errors: list[BaseException] = [] + for cache in self.owned_caches: + crop = getattr(cache, "crop", None) + if callable(crop): + try: + crop(0) + except BaseException as exc: + errors.append(exc) + if self.owner is not None: + current = getattr(self.owner, "_cached_past_key_values", None) + if current is self.target_cache: + setattr(self.owner, "_cached_past_key_values", None) + self.released = True + if errors: + raise errors[0] + + +@dataclass +class PagedMLAExecutionContext: + """No-swap execution context for one engine-owned MLA sequence.""" + + target_cache: PagedCacheAdapter + draft_cache: object | None + mode: str = EXECUTION_CONTEXT_PAGED_MLA + released: bool = False + _owned_draft_caches: list[object] = field(default_factory=list) + + def __post_init__(self) -> None: + self.refresh(self.target_cache, self.draft_cache) + + def refresh( + self, target_cache: object | None, draft_cache: object | None + ) -> None: + if target_cache is not self.target_cache: + raise RuntimeError("paged MLA target cache handle was replaced") + active = [] if draft_cache is None else [draft_cache] + for cache in tuple(self._owned_draft_caches): + if cache in active: + continue + crop = getattr(cache, "crop", None) + if callable(crop): + crop(0) + self._owned_draft_caches.remove(cache) + if ( + draft_cache is not None + and draft_cache not in self._owned_draft_caches + ): + self._owned_draft_caches.append(draft_cache) + self.draft_cache = draft_cache + + @contextmanager + def activate(self) -> Iterator[None]: + if self.released: + raise RuntimeError("paged MLA execution context has been released") + yield + + def release(self) -> None: + if self.released: + return + errors: list[BaseException] = [] + for cache in self._owned_draft_caches: + crop = getattr(cache, "crop", None) + if callable(crop): + try: + crop(0) + except BaseException as exc: + errors.append(exc) + try: + self.target_cache.release() + except BaseException as exc: + errors.append(exc) + self.released = True + if errors: + raise errors[0] + + +@dataclass +class ServingSpecSession: + """Persistent serving record for exactly one request sequence.""" + + spec_session: SpecSession | object + request_id: str + seq_id: int + execution_context: TemporaryDynamicCacheContext | PagedMLAExecutionContext + decode_state: SpecDecodeState + callbacks: tuple[Callable[[object], None], ...] = () + cancelled: bool = False + in_flight: bool = False + pending_draft: object | None = None + streamed_count: int = 0 + output_token_ids: list[int] = field(default_factory=list) + released: bool = False + failure_reason: dict[str, str] | None = None + paged_mla_admission: dict[str, object] = field(default_factory=dict) + paged_mla_block_budget: int = 0 + + @property + def finished(self) -> bool: + return bool(getattr(self.spec_session, "finished", False)) + + def diagnostics(self) -> dict[str, object]: + return { + "request_id": self.request_id, + "seq_id": self.seq_id, + "execution_context": self.execution_context.mode, + "cancelled": self.cancelled, + "in_flight": self.in_flight, + "released": self.released, + "failure_reason": self.failure_reason, + "paged_mla_admission": dict(self.paged_mla_admission), + "cached_len": self.decode_state.cached_len, + "emitted_len": self.decode_state.emitted_len, + } + + +class SpecSessionDriver: + """Step canonical speculative sessions without taking scheduler ownership.""" + + def __init__( + self, + speculator: object, + *, + enable_paged_mla: bool = False, + max_resident_paged_speculative_sessions: int = 1, + min_free_mla_blocks_after_admission: int = 1, + ) -> None: + if ( + type(max_resident_paged_speculative_sessions) is not int + or max_resident_paged_speculative_sessions < 0 + ): + raise ValueError( + "max_resident_paged_speculative_sessions must be an integer >= 0" + ) + if ( + type(min_free_mla_blocks_after_admission) is not int + or min_free_mla_blocks_after_admission < 1 + ): + raise ValueError( + "min_free_mla_blocks_after_admission must be an integer >= 1" + ) + self.speculator = speculator + self.sessions: dict[int, ServingSpecSession] = {} + self.enable_paged_mla = enable_paged_mla + self.max_resident_paged_speculative_sessions = ( + max_resident_paged_speculative_sessions + ) + self.min_free_mla_blocks_after_admission = ( + min_free_mla_blocks_after_admission + ) + self._admission_counters = { + "admitted": 0, + "session_cap": 0, + "free_block_reserve": 0, + "ineligible": 0, + "begin_failed": 0, + } + self._last_admission_rejection: dict[str, object] | None = None + begin_session = getattr(speculator, "begin_session") + begin_parameters = signature(begin_session).parameters.values() + self.supports_request_generator = any( + parameter.kind is Parameter.VAR_KEYWORD + or parameter.name == "generator" + for parameter in begin_parameters + ) + self.supports_target_cache_adapter = any( + parameter.kind is Parameter.VAR_KEYWORD + or parameter.name == "target_cache_adapter" + for parameter in begin_parameters + ) + + @property + def execution_context_mode(self) -> str: + modes = { + record.execution_context.mode for record in self.sessions.values() + } + if not modes: + return EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + if len(modes) == 1: + return next(iter(modes)) + return "mixed" + + @property + def admission_stats(self) -> dict[str, object]: + owner = getattr(self.speculator, "moe", None) + cache = getattr(owner, "_native_mla_cache", None) + return { + "active_sessions": self._active_paged_mla_sessions(), + "enabled": self.enable_paged_mla, + "max_resident_sessions": self.max_resident_paged_speculative_sessions, + "min_free_blocks_after_admission": ( + self.min_free_mla_blocks_after_admission + ), + "free_blocks": ( + cache.free_block_count + if isinstance(cache, MLAPagedKVCache) + else None + ), + "counters": dict(self._admission_counters), + "last_rejection": ( + dict(self._last_admission_rejection) + if self._last_admission_rejection is not None + else None + ), + } + + def begin( + self, + *, + request_id: str, + seq_id: int, + prompt_token_ids: Sequence[int], + max_new_tokens: int, + temperature: float, + top_k: int, + top_p: float, + stop_token_ids: Sequence[int], + callbacks: Sequence[Callable[[object], None]], + generator: torch.Generator | None = None, + ) -> ServingSpecSession: + if seq_id in self.sessions: + raise ValueError( + f"speculative session already exists for seq_id {seq_id}" + ) + begin_session = getattr(self.speculator, "begin_session") + owner = getattr(self.speculator, "moe", None) + admission = self._paged_mla_admission( + owner, temperature, len(prompt_token_ids), max_new_tokens + ) + use_paged_mla = bool(admission["admitted"]) + if owner is not None and not use_paged_mla: + setattr(owner, "_cached_past_key_values", None) + begin_kwargs: dict[str, object] = { + "max_new_tokens": max_new_tokens, + "temperature": temperature, + "stop_token_ids": list(stop_token_ids), + "top_k": top_k, + "top_p": top_p, + "collect_route_union": True, + } + if self.supports_request_generator: + begin_kwargs["generator"] = generator + paged_adapter: PagedCacheAdapter | None = None + try: + paged_block_budget = 0 + if use_paged_mla: + cache = getattr(owner, "_native_mla_cache") + paged_block_budget = self._peak_blocks_for_request( + len(prompt_token_ids), max_new_tokens, cache.block_size + ) + paged_adapter = PagedCacheAdapter( + cache, seq_id=seq_id, initial_length=len(prompt_token_ids) + ) + begin_kwargs["target_cache_adapter"] = paged_adapter + session = begin_session( + torch.tensor([list(prompt_token_ids)], dtype=torch.long), + **begin_kwargs, + ) + if paged_adapter is None: + context: ( + TemporaryDynamicCacheContext | PagedMLAExecutionContext + ) = TemporaryDynamicCacheContext( + owner=owner, + target_cache=getattr(session, "target_kv", None), + draft_cache=getattr(session, "draft_kv", None), + ) + else: + if getattr(session, "target_kv", None) is not paged_adapter: + raise RuntimeError( + "paged MLA session did not retain its cache adapter" + ) + context = PagedMLAExecutionContext( + target_cache=paged_adapter, + draft_cache=getattr(session, "draft_kv", None), + ) + record = ServingSpecSession( + spec_session=session, + request_id=request_id, + seq_id=seq_id, + execution_context=context, + decode_state=SpecDecodeState( + seq_id=seq_id, prompt_len=len(prompt_token_ids) + ), + callbacks=tuple(callbacks), + paged_mla_admission=admission, + paged_mla_block_budget=paged_block_budget, + ) + except BaseException: + if paged_adapter is not None: + paged_adapter.release() + failed_decision = { + "eligible": bool(admission["eligible"]), + "admitted": False, + "reason": "begin_failed", + } + self._last_admission_rejection = failed_decision + self._admission_counters["begin_failed"] += 1 + raise + self.sessions[seq_id] = record + reason = str(admission["reason"]) + self._admission_counters[reason] += 1 + if not bool(admission["admitted"]): + self._last_admission_rejection = dict(admission) + return record + + def draft(self, record: ServingSpecSession) -> object: + self._ensure_runnable(record) + if record.pending_draft is not None: + return record.pending_draft + self._refresh_context(record) + record.in_flight = True + primary: BaseException | None = None + try: + with record.execution_context.activate(): + draft = getattr(self.speculator, "draft_round")( + record.spec_session + ) + record.pending_draft = draft + return draft + except BaseException as exc: + primary = exc + raise + finally: + self._finish_inflight(record, primary) + + def verify(self, record: ServingSpecSession) -> object | None: + self._ensure_runnable(record) + if record.pending_draft is None: + raise RuntimeError("verify requires a pending draft") + self._refresh_context(record) + record.in_flight = True + primary: BaseException | None = None + try: + with record.execution_context.activate(): + return getattr(self.speculator, "verify_round")( + record.spec_session + ) + except BaseException as exc: + primary = exc + raise + finally: + record.pending_draft = None + self._finish_inflight(record, primary) + + def commit(self, record: ServingSpecSession) -> tuple[int, ...]: + """Publish unseen canonical emissions and advance logical accounting.""" + if record.cancelled or record.released: + return () + emitted_obj = getattr(record.spec_session, "emitted", None) + if emitted_obj is None: + emitted_obj = getattr(record.spec_session, "output_ids") + emitted = tuple( + int(token) for token in cast(Sequence[int], emitted_obj) + ) + committed = emitted[record.streamed_count :] + if not committed: + return () + record.streamed_count += len(committed) + record.output_token_ids.extend(committed) + record.decode_state.record_commit(len(committed)) + if not record.decode_state.invariant_holds(): + raise RuntimeError( + "speculative serving logical cache invariant violated for " + f"seq_id {record.seq_id}" + ) + return committed + + def cancel(self, seq_id: int) -> None: + record = self.sessions.get(seq_id) + if record is None: + return + record.cancelled = True + if not record.in_flight: + self.release(record) + + def release(self, record: ServingSpecSession) -> None: + if record.released: + return + errors: list[BaseException] = [] + clear_pending = getattr(record.spec_session, "clear_pending", None) + if callable(clear_pending): + try: + clear_pending() + except BaseException as exc: + errors.append(exc) + record.pending_draft = None + try: + self._refresh_context(record) + except BaseException as exc: + errors.append(exc) + try: + try: + record.execution_context.release() + except BaseException as exc: + errors.append(exc) + finally: + record.released = True + self.sessions.pop(record.seq_id, None) + if errors: + for cleanup_error in errors[1:]: + add_note = getattr(errors[0], "add_note", None) + if callable(add_note): + add_note(f"additional cleanup failure: {cleanup_error}") + raise errors[0] + + def fail( + self, record: ServingSpecSession, failure_reason: dict[str, str] + ) -> None: + record.failure_reason = dict(failure_reason) + record.cancelled = True + self.release(record) + + def _release_if_cancelled(self, record: ServingSpecSession) -> None: + if record.cancelled and not record.released: + self.release(record) + + def _finish_inflight( + self, + record: ServingSpecSession, + primary: BaseException | None, + ) -> None: + cleanup_errors: list[BaseException] = [] + try: + self._refresh_context(record) + except BaseException as exc: + cleanup_errors.append(exc) + record.in_flight = False + try: + self._release_if_cancelled(record) + except BaseException as exc: + cleanup_errors.append(exc) + if primary is not None: + add_note = getattr(primary, "add_note", None) + if callable(add_note): + for cleanup_error in cleanup_errors: + add_note(f"session cleanup failed: {cleanup_error}") + return + if cleanup_errors: + raise cleanup_errors[0] + + @staticmethod + def _refresh_context(record: ServingSpecSession) -> None: + record.execution_context.refresh( + getattr(record.spec_session, "target_kv", None), + getattr(record.spec_session, "draft_kv", None), + ) + + @staticmethod + def _ensure_runnable(record: ServingSpecSession) -> None: + if record.cancelled: + raise RuntimeError("speculative serving session is cancelled") + if record.released: + raise RuntimeError("speculative serving session is released") + if record.finished: + raise RuntimeError("speculative serving session is finished") + + def _paged_mla_admission( + self, + owner: object | None, + temperature: float, + prompt_tokens: int, + max_new_tokens: int, + ) -> dict[str, object]: + if ( + not self.enable_paged_mla + or owner is None + or temperature != 0.0 + or not self.supports_target_cache_adapter + ): + return { + "eligible": False, + "admitted": False, + "reason": "ineligible", + } + cache = getattr(owner, "_native_mla_cache", None) + modules_fn = getattr(owner, "_get_mla_attention_modules", None) + if not ( + isinstance(cache, MLAPagedKVCache) + and callable(modules_fn) + and bool(modules_fn()) + ): + return { + "eligible": False, + "admitted": False, + "reason": "ineligible", + } + if self._transient_verify_tokens() is None: + return { + "eligible": False, + "admitted": False, + "reason": "ineligible", + } + if ( + self._active_paged_mla_sessions() + >= self.max_resident_paged_speculative_sessions + ): + return { + "eligible": True, + "admitted": False, + "reason": "session_cap", + } + declared_blocks = self._peak_blocks_for_request( + prompt_tokens, max_new_tokens, cache.block_size + ) + available_after_active_headroom = ( + cache.free_block_count - self._active_unallocated_headroom(cache) + ) + if ( + available_after_active_headroom - declared_blocks + < self.min_free_mla_blocks_after_admission + ): + return { + "eligible": True, + "admitted": False, + "reason": "free_block_reserve", + } + return {"eligible": True, "admitted": True, "reason": "admitted"} + + @staticmethod + def _blocks_for_tokens(token_count: int, block_size: int) -> int: + return (token_count + block_size - 1) // block_size + + def _transient_verify_tokens(self) -> int | None: + config = getattr(self.speculator, "config", None) + block_size = getattr(config, "block_size", None) + if type(block_size) is not int or block_size < 2: + return None + return block_size - 1 + + def _peak_blocks_for_request( + self, prompt_tokens: int, max_new_tokens: int, mla_block_size: int + ) -> int: + transient_tokens = self._transient_verify_tokens() + if transient_tokens is None: + raise RuntimeError("paged MLA requires a valid DFlash block_size") + return self._blocks_for_tokens( + prompt_tokens + max_new_tokens + transient_tokens, + mla_block_size, + ) + + def _active_unallocated_headroom(self, cache: MLAPagedKVCache) -> int: + headroom = 0 + for record in self.sessions.values(): + if ( + record.released + or record.execution_context.mode != EXECUTION_CONTEXT_PAGED_MLA + ): + continue + allocated = len(cache.get_block_table(record.seq_id)) + headroom += max(0, record.paged_mla_block_budget - allocated) + return headroom + + def _active_paged_mla_sessions(self) -> int: + return sum( + record.execution_context.mode == EXECUTION_CONTEXT_PAGED_MLA + for record in self.sessions.values() + if not record.released + ) + + +__all__ = [ + "EXECUTION_CONTEXT_TEMPORARY_DYNAMIC", + "PagedMLAExecutionContext", + "ServingSpecSession", + "SpecSessionDriver", + "TemporaryDynamicCacheContext", +] diff --git a/moe_infinity/serving/spec_state.py b/moe_infinity/serving/spec_state.py index a94d5970..e130765c 100644 --- a/moe_infinity/serving/spec_state.py +++ b/moe_infinity/serving/spec_state.py @@ -15,9 +15,10 @@ class VerifyStepAccounting: class SpecDecodeState: """Per-sequence cached-vs-emitted bookkeeping for serving-path DFlash. - ``cached_len`` is the number of tokens whose KV is committed to the paged - cache; ``emitted_len`` is the number of committed output tokens returned to - the user. Invariant after every committed step: + ``cached_len`` is the serving-visible logical committed length and + ``emitted_len`` is the number of committed output tokens returned to the + user. In Stage 4a the physical KV remains in a private ``DynamicCache``; + this state does not claim paged-cache ownership. Invariant after every step: ``cached_len == prompt_len + emitted_len``. """ @@ -40,7 +41,8 @@ def record_verify( The verify forward transiently writes KV for all ``block_len`` block tokens; ``committed`` (1..block_len) of them are kept. Returns the target length to which the paged cache must be truncated to drop the rejected - tail, and advances the cached/emitted counters. + tail in the active execution context, and advances the logical + cached/emitted counters. """ if block_len < 1: raise ValueError(f"block_len must be >= 1, got {block_len}") @@ -48,6 +50,27 @@ def record_verify( raise ValueError( f"committed must be in [1, {block_len}], got {committed}" ) + return self.record_commit(committed, block_len=block_len) + + def record_commit( + self, committed: int, *, block_len: int | None = None + ) -> VerifyStepAccounting: + """Advance logical serving counters for one published commit. + + Stage 4a deliberately keeps the physical target KV in a private + ``DynamicCache``. This method records the serving-visible logical + commit only; it does not claim ownership of, append to, or truncate the + paged cache. ``block_len`` is the transient verify width when known and + defaults to the number of committed tokens for the prefill anchor. + """ + if committed < 1: + raise ValueError(f"committed must be >= 1, got {committed}") + transient = committed if block_len is None else block_len + if transient < committed: + raise ValueError( + "block_len must be >= committed; " + f"got block_len={transient}, committed={committed}" + ) self.cached_len += committed self.emitted_len += committed return VerifyStepAccounting( diff --git a/moe_infinity/spec_decode/__init__.py b/moe_infinity/spec_decode/__init__.py index a7310874..66898ebc 100644 --- a/moe_infinity/spec_decode/__init__.py +++ b/moe_infinity/spec_decode/__init__.py @@ -1,6 +1,18 @@ +from moe_infinity.spec_decode.backends import ( + DFlashExecutionBackend, + ExecutionBackend, + PhysicalCohortBackend, + PhysicalCohortResult, +) +from moe_infinity.spec_decode.backends_bare_hf import ( + BareHFCohortResult, + BatchedBareHFBackend, +) +from moe_infinity.spec_decode.backends_rich import BatchedRichBackend from moe_infinity.spec_decode.dflash import ( DFlashConfig, DFlashSpeculator, + build_pairing_evidence, read_dflash_config, validate_pairing, ) @@ -9,13 +21,67 @@ glm_dflash_drafter_for, validate_glm_pairing, ) +from moe_infinity.spec_decode.protocols import ( + BackendCapabilities, + CacheAdapter, + CacheKind, + CacheSnapshot, + DenseCacheAdapter, + ExecutorEvidence, + NativeStepTrace, + PairingEvidence, + RequestSpec, + RichBatchMetadata, + RichForwardResult, + SamplingContext, + SessionRoundResult, + SessionTrace, +) +from moe_infinity.spec_decode.session_driver import ( + BackendProgressError, + CohortPlan, + DriverResult, + PhysicalCohortDriverResult, + SessionCleanupError, + SessionDriver, + UnsupportedRequestError, +) __all__ = [ "DFlashConfig", + "DFlashExecutionBackend", "DFlashSpeculator", + "ExecutionBackend", + "PhysicalCohortBackend", + "PhysicalCohortResult", + "BareHFCohortResult", + "BatchedBareHFBackend", + "BatchedRichBackend", + "build_pairing_evidence", "read_dflash_config", "validate_pairing", "glm_dflash_available", "glm_dflash_drafter_for", "validate_glm_pairing", + "BackendCapabilities", + "CacheAdapter", + "CacheKind", + "CacheSnapshot", + "DenseCacheAdapter", + "ExecutorEvidence", + "NativeStepTrace", + "PairingEvidence", + "RequestSpec", + "RichBatchMetadata", + "RichForwardResult", + "SamplingContext", + "SessionRoundResult", + "SessionTrace", + "CohortPlan", + "BackendProgressError", + "DriverResult", + "PhysicalCohortDriverResult", + "SessionDriver", + "SessionCleanupError", + "UnsupportedRequestError", ] diff --git a/moe_infinity/spec_decode/_dflash_sample_ops.py b/moe_infinity/spec_decode/_dflash_sample_ops.py index 276355b2..5d9cc125 100644 --- a/moe_infinity/spec_decode/_dflash_sample_ops.py +++ b/moe_infinity/spec_decode/_dflash_sample_ops.py @@ -58,6 +58,27 @@ class SampledAcceptance(NamedTuple): final_token: int # residual correction (reject) or bonus (full accept) +def _validate_generator_device( + generator: Optional[torch.Generator], probability_device: torch.device +) -> None: + """Require an explicit request generator on the probability tensor device.""" + if generator is None: + return + generator_device = torch.device(generator.device) + probability_device = torch.device(probability_device) + same_device_type = generator_device.type == probability_device.type + compatible_index = ( + generator_device.index is None + or probability_device.index is None + or generator_device.index == probability_device.index + ) + if not (same_device_type and compatible_index): + raise ValueError( + f"generator device {generator_device} does not match " + f"probability device {probability_device}" + ) + + def warped_probs( logits: torch.Tensor, temperature: float = 1.0, @@ -136,13 +157,19 @@ def acceptance_sampled( for seeded determinism; ``None`` uses the global torch RNG (seed it with ``torch.manual_seed``). """ + _validate_generator_device(generator, target_probs.device) num_drafts = int(drafts.shape[0]) for i in range(num_drafts): token = int(drafts[i]) q = float(draft_probs[i, token]) p = float(target_probs[i, token]) accept_prob = min(1.0, p / q) if q > 0 else 0.0 - if float(torch.rand((), generator=generator)) < accept_prob: + if ( + float( + torch.rand((), device=target_probs.device, generator=generator) + ) + < accept_prob + ): continue correction = torch.multinomial( residual_distribution(target_probs[i], draft_probs[i]), diff --git a/moe_infinity/spec_decode/_route_ahead_ctx.py b/moe_infinity/spec_decode/_route_ahead_ctx.py index b99e0767..2a32c2ef 100644 --- a/moe_infinity/spec_decode/_route_ahead_ctx.py +++ b/moe_infinity/spec_decode/_route_ahead_ctx.py @@ -33,6 +33,9 @@ route_ahead_stats: contextvars.ContextVar[Optional[Any]] = ( contextvars.ContextVar("dflash_route_ahead_stats", default=None) ) +route_ahead_row_offsets: contextvars.ContextVar[tuple[int, ...]] = ( + contextvars.ContextVar("dflash_route_ahead_row_offsets", default=()) +) def is_active() -> bool: @@ -47,9 +50,15 @@ def current_stats() -> Optional[Any]: return route_ahead_stats.get() +def current_row_offsets() -> tuple[int, ...]: + return route_ahead_row_offsets.get() + + @contextmanager def route_ahead_context( - prefetcher: Optional[Any] = None, stats: Optional[Any] = None + prefetcher: Optional[Any] = None, + stats: Optional[Any] = None, + row_offsets: tuple[int, ...] = (), ) -> Iterator[None]: """Activate the route-ahead context; token-reset in ``finally``. @@ -61,9 +70,11 @@ def route_ahead_context( active_token = route_ahead_active.set(True) prefetcher_token = route_ahead_prefetcher.set(prefetcher) stats_token = route_ahead_stats.set(stats) + offsets_token = route_ahead_row_offsets.set(tuple(row_offsets)) try: yield finally: + route_ahead_row_offsets.reset(offsets_token) route_ahead_stats.reset(stats_token) route_ahead_prefetcher.reset(prefetcher_token) route_ahead_active.reset(active_token) @@ -76,5 +87,6 @@ def route_ahead_context( "is_active", "current_prefetcher", "current_stats", + "current_row_offsets", "route_ahead_context", ] diff --git a/moe_infinity/spec_decode/_route_ahead_stats.py b/moe_infinity/spec_decode/_route_ahead_stats.py index 271dccbc..dc262489 100644 --- a/moe_infinity/spec_decode/_route_ahead_stats.py +++ b/moe_infinity/spec_decode/_route_ahead_stats.py @@ -47,6 +47,7 @@ rejected_expert_ids, union_experts_from_mask, ) +from moe_infinity.spec_decode.protocols import ExecutorEvidence class RouteAheadStepSummary(NamedTuple): @@ -98,6 +99,14 @@ def __init__(self) -> None: self.kept_prefetch_bytes: int = 0 self.wasted_prefetch_bytes: int = 0 self._bytes_seen: bool = False + self._attempted_layers: List[int] = [] + self._fired_layers: List[int] = [] + self._actual_expert_union: set[tuple[int, int]] = set() + self._actual_expert_union_by_row: set[tuple[int, int, int]] = set() + self._prefetcher_present: bool = False + self._attempted_prefetch_bytes: int = 0 + self._cache_hit_rate: Optional[float] = None + self._fallback_reason: Optional[str] = None self._pending: List[ Tuple[int, List[int], torch.Tensor, Optional[Dict[int, int]]] ] = [] @@ -153,6 +162,38 @@ def observe_layer( (int(layer_id), [int(e) for e in predicted_ids], mask_cpu, nbytes) ) + def observe_executor_attempt( + self, + layer_id: int, + actual_ids: Sequence[int], + *, + actual_ids_by_row: Sequence[tuple[int, int, int]] = (), + prefetcher_present: bool, + fired: bool, + fallback_reason: Optional[str] = None, + prefetched_bytes: int = 0, + cache_hit_rate: Optional[float] = None, + ) -> None: + """Record executor capability/firing without affecting dispatch.""" + layer = int(layer_id) + self._attempted_layers.append(layer) + self._actual_expert_union.update( + (layer, int(expert_id)) for expert_id in actual_ids + ) + self._actual_expert_union_by_row.update(actual_ids_by_row) + self._prefetcher_present = self._prefetcher_present or bool( + prefetcher_present + ) + if fired: + self._fired_layers.append(layer) + if fallback_reason is not None and self._fallback_reason is None: + self._fallback_reason = fallback_reason + self._attempted_prefetch_bytes += max(0, int(prefetched_bytes)) + if cache_hit_rate is not None: + rate = float(cache_hit_rate) + if 0.0 <= rate <= 1.0: + self._cache_hit_rate = rate + def commit_step(self, kept_rows: int) -> RouteAheadStepSummary: """Finalize the in-flight step: coverage + rejected-token waste. @@ -243,11 +284,33 @@ def waste_ratio(self) -> float: return 0.0 return self.wasted_experts / self.predicted_experts + @property + def executor_evidence(self) -> ExecutorEvidence: + """Immutable snapshot, separate from target/drafter pairing.""" + attempted = tuple(self._attempted_layers) + return ExecutorEvidence( + wiring_reachable=bool(attempted), + prefetcher_present=self._prefetcher_present, + attempted_layers=attempted, + fired_layers=tuple(self._fired_layers), + actual_expert_union=frozenset(self._actual_expert_union), + actual_expert_union_by_row=frozenset( + self._actual_expert_union_by_row + ), + prefetched_bytes=self._attempted_prefetch_bytes, + coverage=self.coverage if attempted else None, + wasted_prefetch_bytes=( + self.wasted_prefetch_bytes if self._bytes_seen else None + ), + cache_hit_rate=self._cache_hit_rate, + fallback_reason=self._fallback_reason, + ) + def reset(self) -> None: """Zero all counters and drop any uncommitted records.""" self.__init__() - def as_dict(self) -> Dict[str, Union[int, float, None]]: + def as_dict(self) -> Dict[str, object]: """Flat snapshot of the counters, byte totals, and derived ratios. The three ``*_prefetch_bytes`` entries are ``None`` until a step is @@ -273,6 +336,7 @@ def as_dict(self) -> Dict[str, Union[int, float, None]]: "wasted_prefetch_bytes": ( self.wasted_prefetch_bytes if self._bytes_seen else None ), + "executor_evidence": self.executor_evidence.as_dict(), } diff --git a/moe_infinity/spec_decode/backends.py b/moe_infinity/spec_decode/backends.py new file mode 100644 index 00000000..0c2d8f6c --- /dev/null +++ b/moe_infinity/spec_decode/backends.py @@ -0,0 +1,287 @@ +"""Execution backend contract and singleton DFlash adapter. + +Task 4 deliberately keeps physical model execution per request. Backends may +advertise cohort compatibility now, but ``DFlashExecutionBackend`` always +drives the canonical ``SpecSession`` one row at a time; a later backend can +replace that physical execution without changing the driver contract. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import ( + TYPE_CHECKING, + Hashable, + Protocol, + TypeVar, + runtime_checkable, +) + +import torch + +from moe_infinity.spec_decode.protocols import ( + BackendCapabilities, + ExecutorEvidence, + NativeStepTrace, + PairingEvidence, + RequestSpec, + SamplingContext, + SessionRoundResult, + SessionTrace, +) + +if TYPE_CHECKING: + from moe_infinity.spec_decode.dflash import DFlashSpeculator, SpecSession + + +SessionT = TypeVar("SessionT") +SnapshotT = TypeVar("SnapshotT") + + +@runtime_checkable +class ExecutionBackend(Protocol[SessionT, SnapshotT]): + """Request-scoped speculative execution lifecycle. + + ``supports`` and ``cohort_key`` are pure capability checks: the driver + invokes them for every row before any backend receives ``prefill``. + ``snapshot``/``restore`` form an abort-only boundary. A restored session is + invalid for further execution and must be released without being resumed. + Every successful ``draft`` + ``verify`` pair must increase observable + output/round progress or report the session finished. + """ + + name: str + capabilities: BackendCapabilities + + def supports(self, request: RequestSpec) -> bool: ... + + def cohort_key(self, request: RequestSpec) -> Hashable: ... + + def prefill(self, request: RequestSpec) -> SessionT: ... + + def draft(self, session: SessionT) -> object: ... + + def verify(self, session: SessionT) -> SessionRoundResult: ... + + def snapshot(self, session: SessionT) -> SnapshotT: + """Capture state solely for abort-only restoration.""" + ... + + def restore(self, session: SessionT, snapshot: SnapshotT) -> None: + """Abort to ``snapshot``; the driver must not resume this session.""" + ... + + def is_finished(self, session: SessionT) -> bool: ... + + def output_token_ids(self, session: SessionT) -> tuple[int, ...]: ... + + def trace(self, session: SessionT) -> SessionTrace: ... + + def release(self, session: SessionT) -> None: ... + + +class PhysicalCohortResult(Protocol): + """Validated result surface shared by physical cohort backends.""" + + generated_token_ids: tuple[tuple[int, ...], ...] + step_trace: tuple[NativeStepTrace, ...] + target_cache: object + draft_cache: object + session_traces: tuple[SessionTrace, ...] + + +@runtime_checkable +class PhysicalCohortBackend(Protocol): + """Backend that executes one compatible cohort as a physical batch. + + This contract is deliberately separate from ``ExecutionBackend`` so the + Task 4 request-scoped lifecycle remains valid. Tensor compatibility + adapters can use this seam when dense-cache mechanics require rows to be + drafted, verified, and rolled back together. + """ + + name: str + capabilities: BackendCapabilities + + def supports(self, request: RequestSpec) -> bool: ... + def cohort_key(self, request: RequestSpec) -> Hashable: ... + def execute_cohort( + self, + input_ids: torch.Tensor, + *, + max_new_tokens: tuple[int, ...], + stop_token_ids: tuple[int, ...], + attention_mask: torch.Tensor, + sampling_contexts: tuple[SamplingContext, ...] | None = None, + stop_token_ids_by_row: tuple[tuple[int, ...], ...] | None = None, + ) -> PhysicalCohortResult: ... + + +@dataclass(frozen=True) +class _DFlashSnapshot: + emitted: tuple[int, ...] + finished: bool + round_index: int + trace_length: int + + +class DFlashExecutionBackend: + """Drive canonical ``SpecSession`` objects one request at a time.""" + + name = "dflash-per-request" + + def __init__( + self, + speculator: DFlashSpeculator, + *, + retain_diagnostics: bool = False, + collect_route_union: bool = True, + ) -> None: + self.speculator = speculator + self.retain_diagnostics = retain_diagnostics + self.collect_route_union = collect_route_union + self._restored_sessions: set[int] = set() + rich_forward = callable( + getattr(speculator.moe, "_native_model_forward_rich", None) + ) + pairing_evidence = getattr( + speculator, "pairing_evidence", PairingEvidence() + ) + executor_evidence = getattr( + speculator, "executor_evidence", ExecutorEvidence() + ) + self.capabilities = BackendCapabilities( + supports_batch=False, + supports_sampling=True, + supports_ragged_rows=True, + cache_kind="dense_dynamic", + supports_route_ahead=executor_evidence.wiring_reachable, + supports_rich_forward=rich_forward, + pairing_evidence=pairing_evidence, + executor_evidence=executor_evidence, + ) + + def supports(self, request: RequestSpec) -> bool: + del request + return True + + def cohort_key(self, request: RequestSpec) -> Hashable: + return (request.is_sampled, self.capabilities.cache_kind) + + def prefill(self, request: RequestSpec) -> SpecSession: + sampling = request.sampling + session = self.speculator.begin_session( + torch.tensor([request.prompt_token_ids], dtype=torch.long), + max_new_tokens=request.max_new_tokens, + temperature=sampling.temperature, + stop_token_ids=list(request.stop_token_ids), + top_k=sampling.top_k, + top_p=sampling.top_p, + generator=sampling.generator, + collect_route_union=self.collect_route_union, + ) + if len(session.output_ids) >= request.max_new_tokens: + session.finished = True + return session + + def draft(self, session: SpecSession) -> object: + return self.speculator.draft_round(session) + + def verify(self, session: SpecSession) -> SessionRoundResult: + verified = self.speculator.verify_round(session) + tokens = tuple(int(token) for token in verified.accepted_token_ids) + includes_bonus = len(tokens) == verified.verified_accept + 1 + return SessionRoundResult( + accepted_draft_count=verified.verified_accept, + committed_token_ids=tokens, + next_anchor=tokens[-1] if includes_bonus else None, + target_cache_length=session.start, + emitted_length=len(session.output_ids), + finished=verified.finished, + finish_reason=( + self._finish_reason(session) if verified.finished else None + ), + ) + + def snapshot(self, session: SpecSession) -> object: + return _DFlashSnapshot( + emitted=tuple(session.emitted), + finished=session.finished, + round_index=session.round_index, + trace_length=len(session.step_trace), + ) + + def restore(self, session: SpecSession, snapshot: object) -> None: + if not isinstance(snapshot, _DFlashSnapshot): + raise TypeError("invalid DFlash session snapshot") + session.emitted[:] = snapshot.emitted + session.finished = snapshot.finished + session.round_index = snapshot.round_index + del session.step_trace[snapshot.trace_length :] + session.clear_pending() + self._restored_sessions.add(id(session)) + + def is_finished(self, session: SpecSession) -> bool: + return session.finished + + def output_token_ids(self, session: SpecSession) -> tuple[int, ...]: + return tuple(session.output_ids[: session.max_new_tokens]) + + def trace(self, session: SpecSession) -> SessionTrace: + stats = getattr(self.speculator, "route_ahead_stats", None) + executor_evidence = ( + stats.executor_evidence + if stats is not None and stats.executor_evidence.attempted_layers + else self.capabilities.executor_evidence + ) + trace = SessionTrace( + request_id="", + backend=self.name, + cache_kind=self.capabilities.cache_kind, + sampled=session.sampled, + finish_reason=self._finish_reason(session), + route_ahead_status=( + "enabled" + if self.capabilities.supports_route_ahead + else "disabled" + ), + pairing_evidence=self.capabilities.pairing_evidence, + executor_evidence=executor_evidence, + ) + for step in session.step_trace: + trace.append(step) + trace.finish_reason = self._finish_reason(session) + return trace + + def release(self, session: SpecSession) -> None: + restored = id(session) in self._restored_sessions + self._restored_sessions.discard(id(session)) + if self.retain_diagnostics and not restored: + self.speculator.last_target_cache = session.target_kv + self.speculator.last_draft_cache = session.draft_kv + return + errors: list[BaseException] = [] + for cache in (session.target_kv, session.draft_kv): + crop = getattr(cache, "crop", None) + if callable(crop): + try: + crop(0) + except BaseException as exc: + errors.append(exc) + if errors: + raise errors[0] + + @staticmethod + def _finish_reason(session: SpecSession) -> str: + output = session.output_ids + if output and output[-1] in session.stop_ids: + return "stop" + return "length" + + +__all__ = [ + "DFlashExecutionBackend", + "ExecutionBackend", + "PhysicalCohortBackend", + "PhysicalCohortResult", +] diff --git a/moe_infinity/spec_decode/backends_bare_hf.py b/moe_infinity/spec_decode/backends_bare_hf.py new file mode 100644 index 00000000..e26a9a77 --- /dev/null +++ b/moe_infinity/spec_decode/backends_bare_hf.py @@ -0,0 +1,449 @@ +"""Physical-cohort greedy and sampled execution for bare HF targets.""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Any, Hashable + +import torch + +from moe_infinity.spec_decode._dflash_ops import ( + acceptance_lengths, + build_block_with_prefixes, + committed_tokens_ragged, +) +from moe_infinity.spec_decode._dflash_sample_ops import ( + _validate_generator_device, + acceptance_sampled, + committed_tokens_sampled, + warped_probs, +) +from moe_infinity.spec_decode.protocols import ( + BackendCapabilities, + ExecutorEvidence, + NativeStepTrace, + RequestSpec, + SamplingContext, + SessionTrace, +) + + +@dataclass(frozen=True) +class BareHFCohortResult: + """Generated rows and compatibility diagnostics from one physical cohort.""" + + generated_token_ids: tuple[tuple[int, ...], ...] + step_trace: tuple[NativeStepTrace, ...] + target_cache: Any + draft_cache: Any + session_traces: tuple[SessionTrace, ...] = () + + @property + def generated_lengths(self) -> tuple[int, ...]: + return tuple(len(row) for row in self.generated_token_ids) + + +class BatchedBareHFBackend: + """Execute a mixed greedy/sampled cohort with one shared dense cache. + + Rows use lockstep minimum-commit rollback. A row that commits farther than + the shared physical cache carries its already-emitted tail into the next + verify block as a known prefix; those re-confirmed tokens are not emitted a + second time. + """ + + name = "dflash-batched-bare-hf" + allows_rich_forward = False + + def __init__(self, speculator: Any) -> None: + self.speculator = speculator + pairing_evidence = getattr(speculator, "pairing_evidence", None) + executor_evidence = ExecutorEvidence( + wiring_reachable=False, + prefetcher_present=False, + fallback_reason="executor_unreachable", + ) + self.capabilities = BackendCapabilities( + supports_batch=True, + supports_sampling=True, + supports_ragged_rows=True, + cache_kind="dense_dynamic", + supports_route_ahead=False, + supports_rich_forward=False, + **( + {"pairing_evidence": pairing_evidence} + if pairing_evidence is not None + else {} + ), + executor_evidence=executor_evidence, + ) + + def supports(self, request: RequestSpec) -> bool: + del request + rich = callable( + getattr( + self.speculator.moe, + "_native_model_forward_rich", + None, + ) + ) + return (not rich) or self.allows_rich_forward + + def cohort_key(self, request: RequestSpec) -> Hashable: + del request + return ("mixed-greedy-sampled", self.capabilities.cache_kind) + + @torch.no_grad() + def execute_cohort( + self, + input_ids: torch.Tensor, + *, + max_new_tokens: tuple[int, ...], + stop_token_ids: tuple[int, ...], + attention_mask: torch.Tensor, + sampling_contexts: tuple[SamplingContext, ...] | None = None, + stop_token_ids_by_row: tuple[tuple[int, ...], ...] | None = None, + ) -> BareHFCohortResult: + """Run dense-cache batching while preserving each row's sampler.""" + from transformers import DynamicCache + + spec = self.speculator + if ( + callable(getattr(spec.moe, "_native_model_forward_rich", None)) + and not self.allows_rich_forward + ): + raise NotImplementedError( + "BatchedBareHFBackend requires a bare HF target; the MoE " + "rich-forward seam is batch==1 only" + ) + if input_ids.ndim != 2: + raise ValueError( + "BatchedBareHFBackend expects input_ids of shape [batch, seq], " + f"got {tuple(input_ids.shape)}" + ) + + batch, padded_prompt = int(input_ids.shape[0]), int(input_ids.shape[1]) + if len(max_new_tokens) != batch: + raise ValueError( + f"per-sequence max_new_tokens has {len(max_new_tokens)} entries " + f"for batch size {batch}" + ) + budgets = [int(value) for value in max_new_tokens] + if any(value < 0 for value in budgets): + raise ValueError(f"max_new_tokens must be >= 0, got {budgets}") + + input_ids = input_ids.to(spec.device) + attention_mask = attention_mask.to(device=spec.device) + if tuple(attention_mask.shape) != tuple(input_ids.shape): + raise ValueError( + f"attention_mask shape {tuple(attention_mask.shape)} != " + f"input_ids shape {tuple(input_ids.shape)}" + ) + binary = (attention_mask == 0) | (attention_mask == 1) + if not bool(torch.all(binary).item()): + raise ValueError("attention_mask must be 0/1 valued") + attention_mask = attention_mask.to(dtype=torch.long) + if int(attention_mask[:, -1].min()) != 1: + raise ValueError( + "batched DFlash requires LEFT-padded prompts: every row's last " + "token must be real (attention_mask[:, -1] == 1)" + ) + steps = attention_mask[:, 1:] - attention_mask[:, :-1] + if steps.numel() and int(steps.min()) < 0: + raise ValueError( + "batched DFlash requires LEFT-padded prompts: each " + "attention_mask row must be 0*1* (pads first, then real tokens)" + ) + + if sampling_contexts is None: + samplings = tuple(SamplingContext() for _ in range(batch)) + else: + samplings = tuple(sampling_contexts) + if len(samplings) != batch: + raise ValueError( + f"sampling_contexts has {len(samplings)} entries for " + f"batch size {batch}" + ) + for row, sampling in enumerate(samplings): + if not isinstance(sampling, SamplingContext): + raise TypeError( + f"sampling_contexts[{row}] must be SamplingContext" + ) + if sampling.is_sampled and budgets[row] > 0: + _validate_generator_device(sampling.generator, input_ids.device) + if stop_token_ids_by_row is None: + stop_sets = tuple(set(stop_token_ids) for _ in range(batch)) + else: + if len(stop_token_ids_by_row) != batch: + raise ValueError( + f"stop_token_ids_by_row has {len(stop_token_ids_by_row)} " + f"entries for batch size {batch}" + ) + stop_sets = tuple(set(row) for row in stop_token_ids_by_row) + + block_size = int(spec.config.block_size) + layer_ids = list(spec.config.target_layer_ids) + mask_token_id = int(spec.config.mask_token_id) + pads = padded_prompt - attention_mask.sum(dim=1) + + spec._configure_target_hooks(input_ids) + prefill_position_ids = (attention_mask.cumsum(dim=-1) - 1).clamp_min(0) + logits, hidden_states, target_kv = spec._forward_target( + input_ids, + past_key_values=None, + logits_to_keep=1, + attention_mask=attention_mask, + position_ids=prefill_position_ids, + ) + greedy_anchors = logits[:, -1, :].argmax(dim=-1) + anchors: list[int] = [] + for row, sampling in enumerate(samplings): + if budgets[row] <= 0 or sampling.is_greedy: + anchors.append(int(greedy_anchors[row])) + continue + anchor_probs = warped_probs( + logits[row, -1], + sampling.temperature, + sampling.top_k, + sampling.top_p, + ) + anchors.append( + int( + torch.multinomial( + anchor_probs, + num_samples=1, + generator=sampling.generator, + ) + ) + ) + context_feature = spec._extract_context_feature( + hidden_states, layer_ids + ).to(spec.device) + emitted: list[list[int]] = [ + ([anchors[row]] if budgets[row] > 0 else []) for row in range(batch) + ] + session_traces = [ + SessionTrace( + request_id=f"direct-{row}", + backend=self.name, + cache_kind=self.capabilities.cache_kind, + sampled=samplings[row].is_sampled, + route_ahead_status="disabled", + pairing_evidence=self.capabilities.pairing_evidence, + executor_evidence=self.capabilities.executor_evidence, + ) + for row in range(batch) + ] + finished = [ + budgets[row] <= 0 + or (bool(stop_sets[row]) and anchors[row] in stop_sets[row]) + for row in range(batch) + ] + start = padded_prompt + draft_kv = DynamicCache() if spec._drafter_has_kv_cache else None + step_trace: list[NativeStepTrace] = [] + + def active(row: int) -> bool: + return not finished[row] and len(emitted[row]) < budgets[row] + + while any(active(row) for row in range(batch)): + prev_start = start + pendings = [ + len(emitted[row]) - (start - padded_prompt) + for row in range(batch) + ] + prefixes = [ + emitted[row][start - padded_prompt :] if active(row) else [] + for row in range(batch) + ] + block = build_block_with_prefixes( + prefixes, mask_token_id, block_size + ).to(spec.device) + + drafter_out = spec._run_drafter( + block, context_feature, start, draft_kv + ) + draft_logits = spec.lm_head(drafter_out)[:, -(block_size - 1) :, :] + draft_prob_rows: list[torch.Tensor | None] = [None] * batch + sampled_reconstruction = [False] * batch + for row, sampling in enumerate(samplings): + if not active(row): + continue + pending = pendings[row] + if sampling.is_greedy: + if pending < block_size: + block[row, pending:] = draft_logits[ + row, pending - 1 : + ].argmax(dim=-1) + continue + if pending > 1: + # Tokens beyond the physical shared-cache position were + # already chosen by this row's prior logical verify. Feed + # them only to reconstruct cache state; do not draft, + # accept, emit, or consume request RNG again. + sampled_reconstruction[row] = True + continue + probabilities = warped_probs( + draft_logits[row], + sampling.temperature, + sampling.top_k, + sampling.top_p, + ) + draft_prob_rows[row] = probabilities + for slot in range(block_size - 1): + block[row, slot + 1] = torch.multinomial( + probabilities[slot], + num_samples=1, + generator=sampling.generator, + ) + + cache_snapshot = spec._snapshot_target_cache(target_kv) + block_attention = torch.cat( + [ + attention_mask, + torch.ones( + batch, + start - padded_prompt + block_size, + dtype=attention_mask.dtype, + device=spec.device, + ), + ], + dim=1, + ) + block_position_ids = torch.arange( + start, + start + block_size, + device=spec.device, + dtype=torch.long, + ).unsqueeze(0) - pads.unsqueeze(1) + logits, hidden_states, target_kv = spec._verify_target_block( + block, + target_kv, + attention_mask=block_attention, + position_ids=block_position_ids, + ) + posterior = logits.argmax(dim=-1).to(spec.device) + greedy_accepts = acceptance_lengths(block, posterior) + greedy_committed = committed_tokens_ragged( + block, posterior, greedy_accepts + ) + + step_cc: dict[int, int] = {} + for row in range(batch): + if not active(row): + continue + pending = pendings[row] + sampling = samplings[row] + if sampled_reconstruction[row]: + step_cc[row] = pending - 1 + continue + if sampling.is_sampled: + draft_probs = draft_prob_rows[row] + assert draft_probs is not None + decision = acceptance_sampled( + draft_probs, + warped_probs( + logits[row], + sampling.temperature, + sampling.top_k, + sampling.top_p, + ), + block[row, 1:], + generator=sampling.generator, + ) + accept = decision.accept + committed = committed_tokens_sampled( + block[row : row + 1], + decision.accept, + decision.final_token, + ) + else: + accept = greedy_accepts[row] + committed = greedy_committed[row] + step_tokens = [ + int(token) for token in committed.emitted[0].tolist() + ] + new_tokens = step_tokens[pending - 1 :] + keep = len(new_tokens) + stop = False + stop_ids = stop_sets[row] + if stop_ids: + for index, token in enumerate(new_tokens): + if token in stop_ids: + keep = index + 1 + stop = True + break + remaining = budgets[row] - len(emitted[row]) + if keep > remaining: + keep = remaining + stop = True + emitted[row].extend(new_tokens[:keep]) + step_cc[row] = min(pending - 1 + keep, accept) + 1 + if stop or len(emitted[row]) >= budgets[row]: + finished[row] = True + + continuing = [row for row in step_cc if active(row)] + min_cc = ( + min(step_cc[row] for row in continuing) + if continuing + else min(step_cc.values()) + ) + spec._rollback_target_cache( + target_kv, + cache_snapshot, + prev_start=prev_start, + committed=min_cc, + block_size=block_size, + ) + start = prev_start + min_cc + assert int(target_kv.get_seq_length()) == start + + if spec.route_ahead_stats is not None: + spec.route_ahead_stats.commit_step(kept_rows=min_cc) + + for row, committed_count in step_cc.items(): + step = NativeStepTrace( + prev_start=prev_start, + accept=committed_count - 1, + start=start, + emitted_len=len(emitted[row]), + target_cache_len=int(target_kv.get_seq_length()), + draft_cache_len=( + int(draft_kv.get_seq_length()) + if draft_kv is not None + else None + ), + ) + step_trace.append(step) + session_traces[row].append(step) + if not continuing: + break + + suffix = spec._extract_context_feature(hidden_states, layer_ids).to( + spec.device + )[:, :min_cc, :] + if spec._drafter_has_kv_cache: + context_feature = suffix + else: + context_feature = torch.cat([context_feature, suffix], dim=1) + + generated = tuple( + tuple(emitted[row][: budgets[row]]) for row in range(batch) + ) + for row, trace in enumerate(session_traces): + trace.emitted = len(generated[row]) + trace.finish_reason = ( + "stop" + if generated[row] and generated[row][-1] in stop_sets[row] + else "length" + ) + return BareHFCohortResult( + generated_token_ids=generated, + step_trace=tuple(step_trace), + target_cache=target_kv, + draft_cache=draft_kv, + session_traces=tuple(session_traces), + ) + + +__all__ = ["BareHFCohortResult", "BatchedBareHFBackend"] diff --git a/moe_infinity/spec_decode/backends_rich.py b/moe_infinity/spec_decode/backends_rich.py new file mode 100644 index 00000000..99f527cf --- /dev/null +++ b/moe_infinity/spec_decode/backends_rich.py @@ -0,0 +1,54 @@ +"""Physical DFlash cohorts for explicitly row-aware rich target wrappers.""" + +from __future__ import annotations + +from typing import Any + +from moe_infinity.spec_decode.backends_bare_hf import BatchedBareHFBackend +from moe_infinity.spec_decode.protocols import BackendCapabilities + + +class BatchedRichBackend(BatchedBareHFBackend): + """Reuse dense lockstep mechanics while injecting rich target forwards.""" + + name = "dflash-batched-rich" + allows_rich_forward = True + + def __init__(self, speculator: Any) -> None: + super().__init__(speculator) + moe = speculator.moe + explicit_declaration = getattr(moe, "__dict__", {}).get( + "_supports_native_rich_batch" + ) + if callable(explicit_declaration): + try: + declared = bool(explicit_declaration()) + except Exception: + declared = False + else: + declared = bool(getattr(moe, "_native_rich_batch_capable", False)) + self.wrapper_supported = declared + base = self.capabilities + self.capabilities = BackendCapabilities( + supports_batch=self.wrapper_supported, + supports_sampling=base.supports_sampling, + supports_ragged_rows=base.supports_ragged_rows, + cache_kind=base.cache_kind, + supports_route_ahead=( + self.wrapper_supported + and bool( + getattr( + speculator.executor_evidence, "wiring_reachable", False + ) + ) + ), + supports_rich_forward=self.wrapper_supported, + pairing_evidence=base.pairing_evidence, + executor_evidence=speculator.executor_evidence, + ) + + def supports(self, request: Any) -> bool: + return self.wrapper_supported and super().supports(request) + + +__all__ = ["BatchedRichBackend"] diff --git a/moe_infinity/spec_decode/dflash.py b/moe_infinity/spec_decode/dflash.py index dc670297..a5da171f 100644 --- a/moe_infinity/spec_decode/dflash.py +++ b/moe_infinity/spec_decode/dflash.py @@ -2,29 +2,28 @@ import inspect from dataclasses import dataclass +from numbers import Integral from types import SimpleNamespace from typing import ( TYPE_CHECKING, Any, Callable, List, - NamedTuple, Optional, Sequence, Union, + cast, ) import torch from moe_infinity.spec_decode._dflash_ops import ( acceptance_length, - acceptance_lengths, build_block, - build_block_with_prefixes, committed_tokens, - committed_tokens_ragged, ) from moe_infinity.spec_decode._dflash_sample_ops import ( + _validate_generator_device, acceptance_sampled, committed_tokens_sampled, warped_probs, @@ -33,6 +32,15 @@ from moe_infinity.spec_decode._prefetch_route import union_experts_from_mask from moe_infinity.spec_decode._route_ahead_ctx import route_ahead_context from moe_infinity.spec_decode._route_ahead_stats import RouteAheadStats +from moe_infinity.spec_decode.protocols import ( + CacheAdapter, + ExecutorEvidence, + NativeStepTrace, + PairingEvidence, + RequestSpec, + SamplingContext, + SessionTrace, +) if TYPE_CHECKING: from moe_infinity.engine.generation_loop import GenerationEngine @@ -102,7 +110,9 @@ def read_dflash_config(draft_hf_config: Any) -> DFlashConfig: DFLASH_TARGET_LAYER_IDS = [1, 9, 17, 25, 33] -def validate_pairing(draft_cfg: DFlashConfig, target_hf_config: Any) -> None: +def validate_pairing( + draft_cfg: DFlashConfig, target_hf_config: Any +) -> PairingEvidence: target_text = _get(target_hf_config, "text_config", target_hf_config) t_hidden = int(_get(target_text, "hidden_size")) t_vocab = int(_get(target_text, "vocab_size")) @@ -136,6 +146,16 @@ def validate_pairing(draft_cfg: DFlashConfig, target_hf_config: Any) -> None: f"DFlash target_layer_ids reference target layer {max(draft_cfg.target_layer_ids)} " f"(capture index {highest_capture}) but target has only {t_layers} layers" ) + return PairingEvidence( + valid=True, + config_valid=True, + dimensions_valid=True, + vocab_valid=True, + mask_valid=True, + layers_valid=True, + block_valid=True, + module_valid=None, + ) def validate_drafter_module(draft_model: Any, draft_cfg: DFlashConfig) -> None: @@ -166,6 +186,68 @@ def validate_drafter( return draft_cfg +_PUBLISHED_DFLASH_PAIRS = frozenset( + { + ("openai/gpt-oss-20b", "z-lab/gpt-oss-20b-DFlash"), + ("openai/gpt-oss-120b", "z-lab/gpt-oss-120b-DFlash"), + } +) + + +def _validated_checkpoint_scope( + target_hf_config: Any, + draft_hf_config: Any, + draft_model_path: str | None = None, +) -> tuple[str, ...]: + """Return checkpoint identities only for explicitly published pairs.""" + target_name = _get(target_hf_config, "_name_or_path") + draft_name = draft_model_path or _get(draft_hf_config, "_name_or_path") + pair = (str(target_name), str(draft_name)) + return pair if pair in _PUBLISHED_DFLASH_PAIRS else () + + +def build_pairing_evidence( + draft_model: Any, + target_hf_config: Any, + draft_cfg: DFlashConfig, + *, + validated_checkpoint_scope: tuple[str, ...] = (), +) -> PairingEvidence: + """Run the authoritative structural checks and return their evidence.""" + base = validate_pairing(draft_cfg, target_hf_config) + validate_drafter_module(draft_model, draft_cfg) + return PairingEvidence( + valid=base.valid, + config_valid=base.config_valid, + dimensions_valid=base.dimensions_valid, + vocab_valid=base.vocab_valid, + mask_valid=base.mask_valid, + layers_valid=base.layers_valid, + block_valid=base.block_valid, + module_valid=True, + validated_checkpoint_scope=validated_checkpoint_scope, + ) + + +def executor_wiring_reachable(moe: Any) -> bool: + """Whether a target can reach ``DistributedExpertExecutor`` dispatch.""" + roots = (moe, getattr(moe, "model", None)) + for root in roots: + if root is None: + continue + if getattr(root, "expert_executor", None) is not None: + return True + if getattr(root, "_executor", None) is not None: + return True + modules = getattr(root, "modules", None) + if not callable(modules): + continue + for module in modules(): + if getattr(module, "expert_executor", None) is not None: + return True + return False + + def _resolve_input_embeddings(target: Any) -> Any: getter = getattr(target, "get_input_embeddings", None) if callable(getter): @@ -231,17 +313,199 @@ def _infer_cuda_device(model: Any) -> str: return "cpu" -def _resolve_stop_ids( - target: Any, stop_token_ids: Optional[List[int]] -) -> List[int]: - if stop_token_ids is not None: - return list(stop_token_ids) - eos = _get(_get(target, "config", target), "eos_token_id") - if eos is None: +def _resolve_stop_ids(target: Any, stop_token_ids: Any) -> List[int]: + value = ( + _get(_get(target, "config", target), "eos_token_id") + if stop_token_ids is None + else stop_token_ids + ) + if value is None: return [] - if isinstance(eos, int): - return [eos] - return [int(x) for x in eos] + if isinstance(value, Integral) and not isinstance(value, bool): + return [int(value)] + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + values = tuple(value) + if all( + isinstance(token, Integral) and not isinstance(token, bool) + for token in values + ): + return [int(token) for token in values] + raise ValueError( + "stop_token_ids must be an integer token ID or a sequence of integer token IDs" + ) + + +def _normalize_per_row( + name: str, value: Any, batch: int, convert: Callable[[Any], Any] +) -> tuple[Any, ...]: + """Expand one scalar or validate one explicit value per batch row.""" + if isinstance(value, Sequence) and not isinstance(value, (str, bytes)): + rows = tuple(value) + if len(rows) != batch: + raise ValueError( + f"per-row {name} has {len(rows)} entries for batch size {batch}" + ) + return tuple(convert(row) for row in rows) + return tuple(convert(value) for _ in range(batch)) + + +def _clone_generator(generator: torch.Generator) -> torch.Generator: + clone = torch.Generator(device=generator.device) + clone.set_state(generator.get_state()) + return clone + + +def _normalize_sampling_contexts( + *, + batch: int, + probability_device: Union[str, torch.device], + temperature: Union[float, Sequence[float]], + top_k: Union[int, Sequence[int]], + top_p: Union[float, Sequence[float]], + generator: Union[ + torch.Generator, + Sequence[Optional[torch.Generator]], + None, + ], +) -> tuple[SamplingContext, ...]: + """Normalize scalar-or-per-row sampling inputs before model prefill. + + Numeric scalars broadcast. A scalar generator is an immutable template + for independent row-local clones; an explicit sequence is retained so + each request-owned stream visibly advances. In a physical batch, omitted + generators are replaced by independently seeded generators on the model + device so sampled rows never share the ambient RNG stream. + """ + temperatures = _normalize_per_row("temperature", temperature, batch, float) + top_ks = _normalize_per_row( + "top_k", top_k, batch, lambda value: max(0, int(value)) + ) + top_ps = _normalize_per_row("top_p", top_p, batch, float) + + if isinstance(generator, Sequence): + generators = tuple(generator) + if len(generators) != batch: + raise ValueError( + f"per-row generator has {len(generators)} entries for batch size {batch}" + ) + if any( + item is not None and not isinstance(item, torch.Generator) + for item in generators + ): + raise TypeError( + "per-row generator values must be torch.Generator or None" + ) + elif generator is None: + generators = (None,) * batch + elif isinstance(generator, torch.Generator): + generators = ( + (generator,) + if batch == 1 + else tuple(_clone_generator(generator) for _ in range(batch)) + ) + else: + raise TypeError( + "generator must be torch.Generator, a per-row sequence, or None" + ) + + if batch > 1: + missing_sampled = [ + row + for row in range(batch) + if temperatures[row] > 0 and generators[row] is None + ] + if missing_sampled: + seeds = torch.randint( + 0, + torch.iinfo(torch.int64).max, + (len(missing_sampled),), + dtype=torch.int64, + ).tolist() + mutable_generators = list(generators) + for row, seed in zip(missing_sampled, seeds): + mutable_generators[row] = torch.Generator( + device=torch.device(probability_device) + ).manual_seed(seed) + generators = tuple(mutable_generators) + + sampled_generator_ids = [ + id(generators[row]) + for row in range(batch) + if temperatures[row] > 0 and generators[row] is not None + ] + if len(set(sampled_generator_ids)) != len(sampled_generator_ids): + raise ValueError( + "the same explicit generator object cannot be shared by sampled " + "rows; pass distinct per-row generators or one scalar generator " + "to request cloned streams" + ) + + return tuple( + SamplingContext( + temperature=temperatures[row], + top_k=top_ks[row], + top_p=top_ps[row], + generator=generators[row], + ) + for row in range(batch) + ) + + +def _normalize_stop_rows( + target: Any, + stop_token_ids: Any, + batch: int, +) -> tuple[tuple[tuple[int, ...], ...], bool]: + """Disambiguate a shared flat stop set from nested per-row stop sets.""" + if stop_token_ids is None or ( + isinstance(stop_token_ids, Integral) + and not isinstance(stop_token_ids, bool) + ): + shared = tuple(_resolve_stop_ids(target, stop_token_ids)) + return (shared,) * batch, False + if not isinstance(stop_token_ids, Sequence) or isinstance( + stop_token_ids, (str, bytes) + ): + raise ValueError( + "stop_token_ids must be an integer, a shared sequence of integers, " + "or one sequence per row" + ) + + values = tuple(stop_token_ids) + scalar_rows = tuple( + isinstance(value, Integral) and not isinstance(value, bool) + for value in values + ) + nested_rows = tuple( + value is None + or (isinstance(value, Sequence) and not isinstance(value, (str, bytes))) + for value in values + ) + if all(scalar_rows): + shared = tuple(int(token) for token in values) + return (shared,) * batch, False + if not all(nested_rows): + raise ValueError( + "stop_token_ids cannot mix shared scalar token IDs with per-row sequences" + ) + if len(values) != batch: + raise ValueError( + f"per-row stop_token_ids has {len(values)} entries for batch size {batch}" + ) + default = tuple(_resolve_stop_ids(target, None)) + rows: list[tuple[int, ...]] = [] + for row in values: + if row is None: + rows.append(default) + continue + try: + rows.append(tuple(_resolve_stop_ids(target, row))) + except ValueError as exc: + raise ValueError( + "each per-row stop_token_ids value must be None or a sequence " + "of integer token IDs" + ) from exc + return tuple(rows), True @dataclass(frozen=True) @@ -446,38 +710,14 @@ def rollback_target_cache( ) -class NativeStepTrace(NamedTuple): - """Per-step state accounting for the native DFlash loop (diagnostics). - - The bonus token is emitted but NOT cached, so after every step the target - cache length equals ``start`` while the absolute emitted length is exactly - one ahead -- conflating the two is the bonus-token trap (oracle ruling #3). - """ - - prev_start: int # cached_len at step entry - # accepted drafts actually committed this step, in [0, block_size - 1]; - # smaller than the accept-rule result when the step is truncated by a - # stop token or the max_new_tokens budget (drafts beyond the cut are - # dropped), so ``start == prev_start + accept + 1`` holds in every branch - accept: int - start: int # cached_len after commit == prev_start + accept + 1 - emitted_len: ( - int # generated tokens emitted so far (accepted drafts + bonus) - ) - target_cache_len: int # target_kv.get_seq_length() after crop - draft_cache_len: Optional[ - int - ] # context-KV length after crop (KV drafter only) - - # --------------------------------------------------------------------------- # Single-round control seam (PD-DFlash Task 6 Step 5) # -# ``SpecSession`` + ``begin_session``/``draft_round``/``verify_round`` externalize -# the ``_generate_single`` loop so the serving engine can interpose the 2-D -# verify scheduler between DRAFT and VERIFY, registering each pending verify's -# EXACT token/byte demand. This is ADDITIVE: ``generate()`` / ``_generate_single`` -# are intentionally left byte-identical and the session runs alongside them. +# ``SpecSession`` + ``begin_session``/``draft_round``/``verify_round`` own the +# canonical single-sequence state machine. ``_generate_single`` is only its +# synchronous driver; serving can interpose the 2-D verify scheduler between +# DRAFT and VERIFY while registering each pending verify's EXACT token/byte +# demand without maintaining a second decode algorithm. # --------------------------------------------------------------------------- @@ -581,13 +821,15 @@ class VerifyResult: ``accepted_token_ids`` are the tokens emitted this round (accepted drafts then the bonus, stop/budget-truncated); ``accept`` is the accepted drafts - committed to KV (cache advance minus one); ``committed_count`` is the number - of emitted tokens appended to the sequence; ``finished`` is set once a stop - id or ``max_new_tokens`` ends the session. + committed to KV (cache advance minus one), while ``verified_accept`` is the + untruncated acceptance-rule result. ``committed_count`` is the number of + emitted tokens appended to the sequence; ``finished`` is set once a stop id + or ``max_new_tokens`` ends the session. """ accepted_token_ids: list[int] accept: int + verified_accept: int committed_count: int finished: bool @@ -598,17 +840,13 @@ class SpecSession: ``begin_session`` seeds this from the prefill/anchor forward; ``draft_round`` and ``verify_round`` advance it one DRAFT and one VERIFY round respectively, - reproducing the ``_generate_single`` loop body exactly but under explicit - engine control (Task 6 Step 5). ``generate()`` is left untouched and - byte-identical; this seam runs ALONGSIDE it. + reproducing singleton ``generate()`` under explicit engine control (Task 6 + Step 5). The synchronous singleton path drives this same state machine. """ input_ids: torch.Tensor max_new_tokens: int - temperature: float - top_k: int - top_p: float - sampled: bool + sampling: SamplingContext block_size: int layer_ids: list[int] mask_token_id: int @@ -629,14 +867,42 @@ class SpecSession: _pending: bool = False _pending_block: Optional[torch.Tensor] = None _pending_prev_start: int = 0 - _pending_draft_probs: Optional[torch.Tensor] = None + pending_draft_probs: Optional[torch.Tensor] = None _pending_cache_snapshot: Any = None + @property + def temperature(self) -> float: + return self.sampling.temperature + + @property + def top_k(self) -> int: + return self.sampling.top_k + + @property + def top_p(self) -> float: + return self.sampling.top_p + + @property + def sampled(self) -> bool: + return self.sampling.is_sampled + @property def output_ids(self) -> list[int]: """All tokens emitted so far (anchor ++ per-round commits), capped.""" return list(self.emitted[: self.max_new_tokens]) + @property + def has_pending_draft(self) -> bool: + """Whether this session owns a tentative block awaiting verification.""" + return self._pending + + def clear_pending(self) -> None: + """Discard tentative draft metadata after verification or abort.""" + self._pending = False + self._pending_block = None + self.pending_draft_probs = None + self._pending_cache_snapshot = None + class DFlashSpeculator: def __init__( @@ -661,10 +927,19 @@ def __init__( ) self.config = read_dflash_config(self.draft.config) - validate_drafter(self.draft, self.target.config, draft_cfg=self.config) + checkpoint_scope = _validated_checkpoint_scope( + self.target.config, self.draft.config, draft_model_path + ) + self.pairing_evidence = build_pairing_evidence( + self.draft, + self.target.config, + self.config, + validated_checkpoint_scope=checkpoint_scope, + ) self.embed_tokens, self.lm_head = bind_shared_weights( self.draft, self.target ) + self.executor_evidence = self._base_executor_evidence() self._init_native_runtime() @classmethod @@ -704,11 +979,19 @@ def from_models( if config is None: config = read_dflash_config(_get(self.draft, "config")) self.config = config - validate_pairing(self.config, self.target.config) - validate_drafter_module(self.draft, self.config) + checkpoint_scope = _validated_checkpoint_scope( + self.target.config, _get(self.draft, "config") + ) + self.pairing_evidence = build_pairing_evidence( + self.draft, + self.target.config, + self.config, + validated_checkpoint_scope=checkpoint_scope, + ) self.embed_tokens, self.lm_head = bind_shared_weights( self.draft, self.target ) + self.executor_evidence = self._base_executor_evidence() self._init_native_runtime() return self @@ -738,6 +1021,10 @@ def _init_native_runtime(self) -> None: # are right-padded in the returned rectangle). None on the batch==1 # path, whose output is never padded. self.last_generated_lengths: Optional[List[int]] = None + self.last_session_traces: tuple[SessionTrace, ...] = () + self.last_session_results: tuple[Any, ...] = () + self.rich_forward_batched = False + self.rich_forward_batch_sizes: list[int] = [] # Track A5 metrics handle; None (default) = instrumentation off, zero # overhead. Set via ``enable_route_ahead_stats``. self.route_ahead_stats: Optional[RouteAheadStats] = None @@ -764,6 +1051,33 @@ def _configure_target_hooks(self, input_ids: torch.Tensor) -> None: if callable(eval_fn): eval_fn() + @staticmethod + def _extract_context_feature( + hidden_states: Sequence[torch.Tensor], layer_ids: Sequence[int] + ) -> torch.Tensor: + return _extract_context_feature(hidden_states, layer_ids) + + @staticmethod + def _snapshot_target_cache(target_kv: Any) -> TargetCacheSnapshot: + return snapshot_target_cache(target_kv) + + @staticmethod + def _rollback_target_cache( + target_kv: Any, + snapshot: TargetCacheSnapshot, + *, + prev_start: int, + committed: int, + block_size: int, + ) -> None: + rollback_target_cache( + target_kv, + snapshot, + prev_start=prev_start, + committed=committed, + block_size=block_size, + ) + def _forward_target( self, input_ids: torch.Tensor, @@ -772,6 +1086,7 @@ def _forward_target( *, attention_mask: Optional[torch.Tensor] = None, position_ids: Optional[torch.Tensor] = None, + attention_metadata: Any = None, ) -> tuple[torch.Tensor, Any, Any]: """Rich target forward -> on-device (logits, hidden_states, past_key_values). @@ -784,19 +1099,51 @@ def _forward_target( """ rich = getattr(self.moe, "_native_model_forward_rich", None) if callable(rich): - metadata = ( - None - if past_key_values is None - else SimpleNamespace(is_prefill=False) - ) - token_ids = [int(t) for t in input_ids[0].tolist()] + from moe_infinity.spec_decode.protocols import RichBatchMetadata + + metadata = attention_metadata + if metadata is None and past_key_values is not None: + metadata = SimpleNamespace(is_prefill=False) + if int(input_ids.shape[0]) > 1: + batch = int(input_ids.shape[0]) + query_lengths = ( + tuple( + int(value) + for value in attention_mask.sum(dim=1).tolist() + ) + if attention_mask is not None + and int(attention_mask.shape[1]) == int(input_ids.shape[1]) + else (int(input_ids.shape[1]),) * batch + ) + offsets = [0] + for length in query_lengths: + offsets.append(offsets[-1] + length) + metadata = RichBatchMetadata( + row_offsets=tuple(offsets), + row_lengths=query_lengths, + attention_mask=attention_mask, + position_ids=position_ids, + cache_handles=(past_key_values,) * batch, + is_prefill=past_key_values is None, + ) + token_ids: list[int] | torch.Tensor = input_ids + self.rich_forward_batch_sizes.append(batch) + else: + token_ids = [int(t) for t in input_ids[0].tolist()] result = rich(token_ids, metadata, logits_to_keep=logits_to_keep) - if not isinstance(result, tuple) or len(result) != 3: + from moe_infinity.spec_decode.protocols import RichForwardResult + + if isinstance(result, RichForwardResult): + logits = result.logits + hidden_states = result.hidden_states + past_key_values = result.cache_handle + elif isinstance(result, tuple) and len(result) == 3: + logits, hidden_states, past_key_values = result + else: raise RuntimeError( "_native_model_forward_rich must return " "(logits, hidden_states, past_key_values)" ) - logits, hidden_states, past_key_values = result if not isinstance(logits, torch.Tensor): raise RuntimeError( "native rich forward logits must be a Tensor" @@ -828,6 +1175,20 @@ def _resolve_route_ahead_prefetcher(self) -> Any: engine = getattr(self.moe, "engine", None) return getattr(engine, "expert_prefetcher", None) + def _base_executor_evidence(self) -> ExecutorEvidence: + reachable = executor_wiring_reachable(self.moe) + prefetcher_present = self._resolve_route_ahead_prefetcher() is not None + reason = None + if not reachable: + reason = "executor_unreachable" + elif not prefetcher_present: + reason = "prefetcher_absent" + return ExecutorEvidence( + wiring_reachable=reachable, + prefetcher_present=prefetcher_present, + fallback_reason=reason, + ) + def _verify_target_block( self, block: torch.Tensor, @@ -872,8 +1233,14 @@ def _verify_target_block( if position_ids is not None: kwargs["position_ids"] = position_ids with nvtx_phase("target_verify"): + row_width = int(block.shape[1]) + row_offsets = tuple( + row * row_width for row in range(int(block.shape[0]) + 1) + ) with route_ahead_context( - self._resolve_route_ahead_prefetcher(), stats=stats + self._resolve_route_ahead_prefetcher(), + stats=stats, + row_offsets=row_offsets, ): return self._forward_target( block, past_key_values=target_kv, logits_to_keep=0, **kwargs @@ -925,7 +1292,9 @@ def begin_session( stop_token_ids: Optional[List[int]] = None, top_k: int = 0, top_p: float = 1.0, + generator: Optional[torch.Generator] = None, collect_route_union: bool = False, + target_cache_adapter: CacheAdapter | None = None, ) -> SpecSession: """Prefill + anchor forward; seed a ``SpecSession`` for engine rounds. @@ -950,8 +1319,19 @@ def begin_session( from transformers import DynamicCache - sampled = float(temperature) > 0 + normalized_top_k = max(0, int(top_k)) + sampling = SamplingContext( + temperature=float(temperature), + top_k=normalized_top_k, + top_p=float(top_p), + generator=generator, + ) + sampled = sampling.is_sampled + if target_cache_adapter is not None and sampled: + raise ValueError("paged target cache requires greedy sampling") input_ids = input_ids.to(self.device) + if sampled: + _validate_generator_device(generator, input_ids.device) self._configure_target_hooks(input_ids) num_prompt_tokens = int(input_ids.shape[1]) @@ -959,14 +1339,44 @@ def begin_session( layer_ids = list(self.config.target_layer_ids) max_new_tokens = int(max_new_tokens) - logits, hidden_states, target_kv = self._forward_target( - input_ids, past_key_values=None, logits_to_keep=1 - ) + if target_cache_adapter is None: + logits, hidden_states, target_kv = self._forward_target( + input_ids, past_key_values=None, logits_to_keep=1 + ) + else: + build_metadata = getattr( + target_cache_adapter, "build_attention_metadata", None + ) + if not callable(build_metadata): + raise TypeError( + "paged target cache adapter must build attention metadata" + ) + metadata = build_metadata( + query_length=num_prompt_tokens, is_prefill=True + ) + logits, hidden_states, returned_handle = self._forward_target( + input_ids, + past_key_values=target_cache_adapter, + logits_to_keep=1, + attention_metadata=metadata, + ) + engine_cache = getattr(target_cache_adapter, "cache", None) + if returned_handle is not engine_cache: + raise RuntimeError( + "paged target forward did not return the engine-owned cache" + ) + target_kv = target_cache_adapter if sampled: anchor = int( torch.multinomial( - warped_probs(logits[0, -1], temperature, top_k, top_p), + warped_probs( + logits[0, -1], + sampling.temperature, + sampling.top_k, + sampling.top_p, + ), num_samples=1, + generator=generator, ).item() ) else: @@ -979,10 +1389,7 @@ def begin_session( session = SpecSession( input_ids=input_ids, max_new_tokens=max_new_tokens, - temperature=float(temperature), - top_k=int(top_k), - top_p=float(top_p), - sampled=sampled, + sampling=sampling, block_size=block_size, layer_ids=layer_ids, mask_token_id=int(self.config.mask_token_id), @@ -1020,7 +1427,7 @@ def draft_round(self, session: SpecSession) -> DraftResult: """ if session.finished: raise RuntimeError("draft_round called on a finished session") - if session._pending: + if session.has_pending_draft: raise RuntimeError( "draft_round called with an un-verified pending block; " "call verify_round first" @@ -1045,17 +1452,22 @@ def draft_round(self, session: SpecSession) -> DraftResult: session.top_p, ) block[:, 1:] = torch.multinomial( - draft_probs, num_samples=1 + draft_probs, + num_samples=1, + generator=session.sampling.generator, ).squeeze(-1) else: block[:, 1:] = draft_logits.argmax(dim=-1) session._pending_block = block session._pending_prev_start = prev_start - session._pending_draft_probs = draft_probs - session._pending_cache_snapshot = snapshot_target_cache( - session.target_kv - ) + session.pending_draft_probs = draft_probs + if isinstance(session.target_kv, CacheAdapter): + session._pending_cache_snapshot = session.target_kv.snapshot() + else: + session._pending_cache_snapshot = snapshot_target_cache( + session.target_kv + ) session._pending = True return DraftResult( @@ -1075,7 +1487,7 @@ def verify_round(self, session: SpecSession) -> VerifyResult: union is collected read-only and turned into the EXACT ``expert_nbytes``-summed demand carried to the next ``draft_round``. """ - if not session._pending or session._pending_block is None: + if not session.has_pending_draft or session._pending_block is None: raise RuntimeError( "verify_round called without a pending draft; " "call draft_round first" @@ -1083,8 +1495,13 @@ def verify_round(self, session: SpecSession) -> VerifyResult: block = session._pending_block prev_start = session._pending_prev_start - draft_probs = session._pending_draft_probs + draft_probs = session.pending_draft_probs cache_snapshot = session._pending_cache_snapshot + paged_target = ( + session.target_kv + if isinstance(session.target_kv, CacheAdapter) + else None + ) collector = session.collector stats = ( @@ -1096,13 +1513,44 @@ def verify_round(self, session: SpecSession) -> VerifyResult: stats.begin_step() with nvtx_phase("target_verify"): with route_ahead_context( - self._resolve_route_ahead_prefetcher(), stats=stats + self._resolve_route_ahead_prefetcher(), + stats=stats, + row_offsets=(0, int(block.numel())), ): - logits, hidden_states, session.target_kv = self._forward_target( - block, - past_key_values=session.target_kv, - logits_to_keep=0, - ) + if paged_target is None: + logits, hidden_states, session.target_kv = ( + self._forward_target( + block, + past_key_values=session.target_kv, + logits_to_keep=0, + ) + ) + else: + paged_target.append(session.block_size) + build_metadata = getattr( + paged_target, "build_attention_metadata", None + ) + if not callable(build_metadata): + raise TypeError( + "paged target cache adapter must build attention metadata" + ) + metadata = build_metadata( + query_length=session.block_size, is_prefill=False + ) + logits, hidden_states, returned_handle = ( + self._forward_target( + block, + past_key_values=paged_target, + logits_to_keep=0, + attention_metadata=metadata, + ) + ) + if returned_handle is not getattr( + paged_target, "cache", None + ): + raise RuntimeError( + "paged target forward replaced the engine-owned cache" + ) if session.sampled: assert draft_probs is not None @@ -1115,6 +1563,7 @@ def verify_round(self, session: SpecSession) -> VerifyResult: session.top_p, ), block[0, 1:], + generator=session.sampling.generator, ) accept = decision.accept committed = committed_tokens_sampled( @@ -1134,7 +1583,7 @@ def verify_round(self, session: SpecSession) -> VerifyResult: keep = j + 1 stop = True break - remaining = session.max_new_tokens - len(session.emitted) + remaining = max(0, session.max_new_tokens - len(session.emitted)) if keep > remaining: keep = remaining stop = True @@ -1142,24 +1591,33 @@ def verify_round(self, session: SpecSession) -> VerifyResult: session.emitted.extend(step_tokens[:keep]) cache_committed = min(keep, accept) + 1 session.start = prev_start + cache_committed - rollback_target_cache( - session.target_kv, - cache_snapshot, - prev_start=prev_start, - committed=cache_committed, - block_size=session.block_size, - block=block, - replay=( - ( - lambda prefix, cache: self._forward_target( - prefix, past_key_values=cache, logits_to_keep=0 - )[2] - ) - if cache_snapshot.linear - else None - ), - ) - assert int(session.target_kv.get_seq_length()) == session.start + if paged_target is None: + rollback_target_cache( + session.target_kv, + cache_snapshot, + prev_start=prev_start, + committed=cache_committed, + block_size=session.block_size, + block=block, + replay=( + ( + lambda prefix, cache: self._forward_target( + prefix, past_key_values=cache, logits_to_keep=0 + )[2] + ) + if cache_snapshot.linear + else None + ), + ) + target_cache_len = int(session.target_kv.get_seq_length()) + else: + paged_target.truncate(session.start) + target_cache_len = paged_target.logical_length() + if target_cache_len != session.start: + raise RuntimeError( + "DFlash target cache length invariant violated: " + f"expected {session.start}, got {target_cache_len}" + ) if collector is not None: collector.commit_step(kept_rows=cache_committed) @@ -1178,7 +1636,7 @@ def verify_round(self, session: SpecSession) -> VerifyResult: accept=cache_committed - 1, start=session.start, emitted_len=len(session.emitted), - target_cache_len=int(session.target_kv.get_seq_length()), + target_cache_len=target_cache_len, draft_cache_len=( int(session.draft_kv.get_seq_length()) if session.draft_kv is not None @@ -1201,10 +1659,7 @@ def verify_round(self, session: SpecSession) -> VerifyResult: ) session.anchor = int(committed.bonus[0, 0].item()) - session._pending = False - session._pending_block = None - session._pending_draft_probs = None - session._pending_cache_snapshot = None + session.clear_pending() session.round_index += 1 if finished: @@ -1215,6 +1670,7 @@ def verify_round(self, session: SpecSession) -> VerifyResult: return VerifyResult( accepted_token_ids=step_tokens[:keep], accept=cache_committed - 1, + verified_accept=accept, committed_count=keep, finished=finished, ) @@ -1224,11 +1680,22 @@ def generate( self, input_ids: torch.Tensor, max_new_tokens: Union[int, Sequence[int]] = 256, - temperature: float = 0.0, - stop_token_ids: Optional[List[int]] = None, - top_k: int = 0, - top_p: float = 1.0, + temperature: Union[float, Sequence[float]] = 0.0, + stop_token_ids: Optional[ + Union[ + int, + Sequence[int], + Sequence[Optional[Sequence[int]]], + ] + ] = None, + top_k: Union[int, Sequence[int]] = 0, + top_p: Union[float, Sequence[float]] = 1.0, attention_mask: Optional[torch.Tensor] = None, + generator: Union[ + torch.Generator, + Sequence[Optional[torch.Generator]], + None, + ] = None, ) -> torch.Tensor: """Native DFlash draft->verify->rollback loop (RFC 1.2). @@ -1262,12 +1729,21 @@ def generate( produced KV for block tokens, never for the bonus. Batching (Track C): ``input_ids`` with batch > 1 dispatches to - ``_generate_batched`` -- greedy-only (``temperature`` must be 0), - bare-HF-target-only (the MoE rich-forward seam stays batch==1), with - LEFT-padded prompts described by ``attention_mask`` (omit it when all - prompts share one length) and a scalar or per-sequence - ``max_new_tokens``. The ragged per-row outputs are right-padded into - the returned rectangle; each row's true new-token count is exposed as + ``_generate_batched`` for bare-HF targets and wrappers that explicitly + satisfy the row-aware rich-forward capability. Unsupported rich, + MLA, and hybrid wrappers retain independent per-request sessions. + ``temperature``, ``top_k``, ``top_p``, generators, budgets, and nested + stop-id sets accept one value per row; numeric scalars and flat stop-id + sets retain their shared/broadcast meaning. + For batch > 1, one scalar generator is cloned to the same initial + state for every row, so identical requests have correlated streams. + Callers wanting independent explicit streams should pass one generator + per row or omit the generator. Seed-exact outputs are not guaranteed + across different batch shapes; the batched API guarantees per-request + order/composition invariance for a fixed request and row-local stream. + Prompts are LEFT-padded according to ``attention_mask`` (omit it when + all prompts share one length). Ragged outputs are right-padded in the + returned rectangle and their true new-token counts are exposed as ``self.last_generated_lengths``. ``attention_mask`` is ignored on the batch==1 path. """ @@ -1275,233 +1751,206 @@ def generate( raise ValueError( f"DFlashSpeculator.generate expects input_ids of shape [batch, seq], got {tuple(input_ids.shape)}" ) - if float(temperature) < 0: - raise ValueError( - f"DFlashSpeculator.generate: temperature must be >= 0, got {temperature}" - ) - if input_ids.shape[0] == 1: - budget = max_new_tokens - if isinstance(budget, Sequence): - if len(budget) != 1: - raise ValueError( - f"per-sequence max_new_tokens has {len(budget)} entries " - f"for batch size 1" - ) - budget = budget[0] - return self._generate_single( + batch = int(input_ids.shape[0]) + sampling_contexts = _normalize_sampling_contexts( + batch=batch, + probability_device=self.device, + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generator, + ) + stop_rows, has_per_row_stops = _normalize_stop_rows( + self.target, stop_token_ids, batch + ) + budgets = self._normalize_budgets(max_new_tokens, batch) + rich_forward = callable( + getattr(self.moe, "_native_model_forward_rich", None) + ) + if batch == 1: + self.rich_forward_batched = False + return self._generate_per_request( input_ids, - max_new_tokens=int(budget), - temperature=float(temperature), - stop_token_ids=stop_token_ids, - top_k=top_k, - top_p=top_p, - ) - if float(temperature) > 0: - raise NotImplementedError( - "batched DFlash (batch > 1) is greedy-only for now; " - f"got temperature {temperature}" + budgets=budgets, + stop_rows=stop_rows, + sampling_contexts=sampling_contexts, + attention_mask=attention_mask, ) - if callable(getattr(self.moe, "_native_model_forward_rich", None)): - raise NotImplementedError( - "batched DFlash (batch > 1) requires a bare HF target; the MoE " - "rich-forward seam is batch==1 (engine-gated) only" + if rich_forward: + from moe_infinity.spec_decode.backends_rich import ( + BatchedRichBackend, ) + + if not BatchedRichBackend(self).wrapper_supported: + self.rich_forward_batched = False + return self._generate_per_request( + input_ids, + budgets=budgets, + stop_rows=stop_rows, + sampling_contexts=sampling_contexts, + attention_mask=attention_mask, + ) + self.rich_forward_batched = rich_forward + self.rich_forward_batch_sizes = [] return self._generate_batched( input_ids, - max_new_tokens=max_new_tokens, - stop_token_ids=stop_token_ids, + max_new_tokens=budgets, + stop_token_ids=(None if has_per_row_stops else list(stop_rows[0])), attention_mask=attention_mask, + sampling_contexts=sampling_contexts, + stop_token_ids_by_row=(stop_rows if has_per_row_stops else None), ) + @staticmethod + def _normalize_budgets( + max_new_tokens: Union[int, Sequence[int]], batch: int + ) -> tuple[int, ...]: + if isinstance(max_new_tokens, Sequence): + budgets = tuple(int(value) for value in max_new_tokens) + if len(budgets) != batch: + raise ValueError( + f"per-sequence max_new_tokens has {len(budgets)} entries " + f"for batch size {batch}" + ) + else: + budgets = (int(max_new_tokens),) * batch + if any(value < 0 for value in budgets): + raise ValueError( + f"max_new_tokens must be >= 0, got {list(budgets)}" + ) + return budgets + @torch.no_grad() - def _generate_single( + def _generate_per_request( self, input_ids: torch.Tensor, - max_new_tokens: int, - temperature: float, - stop_token_ids: Optional[List[int]], - top_k: int, - top_p: float, + *, + budgets: tuple[int, ...], + stop_rows: tuple[tuple[int, ...], ...], + sampling_contexts: tuple[SamplingContext, ...], + attention_mask: Optional[torch.Tensor], ) -> torch.Tensor: - """The v1 single-sequence loop (batch==1), byte-identical to pre-Track-C.""" - sampled = float(temperature) > 0 - - from transformers import DynamicCache - - input_ids = input_ids.to(self.device) - self._configure_target_hooks(input_ids) - - num_prompt_tokens = int(input_ids.shape[1]) - block_size = int(self.config.block_size) - layer_ids = list(self.config.target_layer_ids) - max_new_tokens = int(max_new_tokens) - - logits, hidden_states, target_kv = self._forward_target( - input_ids, past_key_values=None, logits_to_keep=1 - ) - if sampled: - anchor = int( - torch.multinomial( - warped_probs(logits[0, -1], temperature, top_k, top_p), - num_samples=1, - ).item() - ) - else: - anchor = int(logits[:, -1, :].argmax(dim=-1).item()) - context_feature = _extract_context_feature(hidden_states, layer_ids).to( - self.device + """Adapt tensor rows to canonical request sessions in row order.""" + from moe_infinity.spec_decode.backends import DFlashExecutionBackend + from moe_infinity.spec_decode.session_driver import ( + SessionDriver, + UnsupportedRequestError, ) - stop_ids = set(_resolve_stop_ids(self.target, stop_token_ids)) - - emitted: List[int] = [anchor] - start = num_prompt_tokens - draft_kv = DynamicCache() if self._drafter_has_kv_cache else None - - self.step_trace = [] - self.last_generated_lengths = None - if stop_ids and anchor in stop_ids and max_new_tokens >= 1: - # The prefill anchor is itself a stop token: emit it and halt - # before any block is drafted (nothing past EOS may be emitted - # or cached; the anchor/bonus is never cached). - self.last_target_cache = target_kv - self.last_draft_cache = draft_kv - new_ids = torch.tensor( - [emitted], dtype=torch.long, device=input_ids.device - ) - return torch.cat([input_ids, new_ids], dim=1) - - while len(emitted) < max_new_tokens: - prev_start = start - block = build_block( - anchor, self.config.mask_token_id, block_size - ).to(self.device) - - drafter_out = self._run_drafter( - block, context_feature, start, draft_kv - ) - draft_logits = self.lm_head(drafter_out)[:, -(block_size - 1) :, :] - draft_probs: Optional[torch.Tensor] = None - if sampled: - # The accept test divides by the drafter's OWN warped slot - # distributions, so drafts must be genuine draws from them -- - # argmax drafts would void the losslessness proof. - draft_probs = warped_probs( - draft_logits[0], temperature, top_k, top_p + batch = int(input_ids.shape[0]) + if attention_mask is not None and batch > 1: + if tuple(attention_mask.shape) != tuple(input_ids.shape): + raise ValueError( + f"attention_mask shape {tuple(attention_mask.shape)} != " + f"input_ids shape {tuple(input_ids.shape)}" + ) + binary = (attention_mask == 0) | (attention_mask == 1) + if not bool(torch.all(binary).item()): + raise ValueError("attention_mask must be 0/1 valued") + + requests: list[RequestSpec] = [] + for row in range(batch): + row_ids = input_ids[row] + if attention_mask is not None and batch > 1: + row_ids = row_ids[attention_mask[row].to(dtype=torch.bool)] + requests.append( + RequestSpec( + request_id=f"direct-{row}", + prompt_token_ids=tuple( + int(token) for token in row_ids.tolist() + ), + max_new_tokens=budgets[row], + stop_token_ids=frozenset(stop_rows[row]), + sampling=sampling_contexts[row], ) - block[:, 1:] = torch.multinomial( - draft_probs, num_samples=1 - ).squeeze(-1) - else: - block[:, 1:] = draft_logits.argmax(dim=-1) - - cache_snapshot = snapshot_target_cache(target_kv) - - logits, hidden_states, target_kv = self._verify_target_block( - block, target_kv ) - if sampled: - assert draft_probs is not None - decision = acceptance_sampled( - draft_probs, - warped_probs(logits[0], temperature, top_k, top_p), - block[0, 1:], - ) - accept = decision.accept - committed = committed_tokens_sampled( - block, decision.accept, decision.final_token + backend = DFlashExecutionBackend( + self, + retain_diagnostics=True, + collect_route_union=False, + ) + for request in requests: + if ( + request.is_sampled + and not backend.capabilities.supports_sampling + ): + raise UnsupportedRequestError( + f"request {request.request_id!r} has no compatible sampled backend" ) - else: - posterior = logits.argmax(dim=-1).to(self.device) - accept = acceptance_length(block, posterior) - committed = committed_tokens(block, posterior, accept) - - # This step's emitted tokens are [d_1 .. d_accept, bonus]; the - # verify forward produced KV only for [anchor, d_1 .. d_accept]. - # Keeping k emitted tokens (stop index k - 1) therefore commits - # min(k, accept) + 1 cached tokens: the anchor plus the first - # min(k, accept) drafts. The bonus is never cached, so a cut at - # the bonus still commits the full accept + 1 block prefix. - step_tokens = [int(t) for t in committed.emitted[0].tolist()] - keep = accept + 1 - stop = False - if stop_ids: - for j, tok in enumerate(step_tokens): - if tok in stop_ids: - keep = j + 1 - stop = True - break - remaining = max_new_tokens - len(emitted) - if keep > remaining: - keep = remaining - stop = True - - emitted.extend(step_tokens[:keep]) - cache_committed = min(keep, accept) + 1 - start = prev_start + cache_committed - rollback_target_cache( - target_kv, - cache_snapshot, - prev_start=prev_start, - committed=cache_committed, - block_size=block_size, - block=block, - replay=( - ( - lambda prefix, cache: self._forward_target( - prefix, - past_key_values=cache, - logits_to_keep=0, - )[2] - ) - if cache_snapshot.linear - else None - ), - ) - assert int(target_kv.get_seq_length()) == start - - if self.route_ahead_stats is not None: - # A5: finalize this verify step's coverage/waste accounting - # with the kept prefix the accept rule just fixed. Read-only. - self.route_ahead_stats.commit_step(kept_rows=cache_committed) - - self.step_trace.append( - NativeStepTrace( - prev_start=prev_start, - accept=cache_committed - 1, - start=start, - emitted_len=len(emitted), - target_cache_len=int(target_kv.get_seq_length()), - draft_cache_len=( - int(draft_kv.get_seq_length()) - if draft_kv is not None - else None - ), + if not backend.supports(request): + mode = "sampled" if request.is_sampled else "greedy" + raise UnsupportedRequestError( + f"request {request.request_id!r} has no compatible {mode} backend" ) - ) - if stop: - break - suffix = _extract_context_feature(hidden_states, layer_ids).to( - self.device - )[:, : accept + 1, :] - if self._drafter_has_kv_cache: - context_feature = suffix - else: - context_feature = torch.cat([context_feature, suffix], dim=1) + # The rich MoE seam stores its dense cache on the shell, so only one + # request may be live at a time. Each row still uses the canonical + # driver/backend lifecycle; this is semantic batching, not a claim of + # physical model batching. + results = tuple( + SessionDriver([backend]).run(request)[0] for request in requests + ) + self.last_session_results = results + self.last_session_traces = tuple(result.trace for result in results) + lengths = [len(result.output_token_ids) for result in results] + self.last_generated_lengths = None if batch == 1 else lengths - anchor = int(committed.bonus[0, 0].item()) + pad_id = _get(_get(self.target, "config", self.target), "pad_token_id") + pad_id = 0 if pad_id is None else int(pad_id) + width = max(lengths, default=0) + new_ids = torch.full( + (batch, width), + pad_id, + dtype=input_ids.dtype, + device=input_ids.device, + ) + for row, result in enumerate(results): + if result.output_token_ids: + new_ids[row, : lengths[row]] = torch.tensor( + result.output_token_ids, + dtype=input_ids.dtype, + device=input_ids.device, + ) + return torch.cat([input_ids, new_ids], dim=1) - self.last_target_cache = target_kv - self.last_draft_cache = draft_kv + @torch.no_grad() + def _generate_single( + self, + input_ids: torch.Tensor, + max_new_tokens: int, + temperature: float, + stop_token_ids: Optional[List[int]], + top_k: int, + top_p: float, + generator: Optional[torch.Generator] = None, + ) -> torch.Tensor: + """Synchronously drive the canonical single-sequence session state.""" + budget = max(0, int(max_new_tokens)) + stop_ids = _resolve_stop_ids(self.target, stop_token_ids) + session = self.begin_session( + input_ids, + max_new_tokens=budget, + temperature=float(temperature), + stop_token_ids=stop_ids, + top_k=int(top_k), + top_p=float(top_p), + generator=generator, + ) + if len(session.output_ids) >= budget: + session.finished = True + while not session.finished: + self.draft_round(session) + self.verify_round(session) + self.last_target_cache = session.target_kv + self.last_draft_cache = session.draft_kv new_ids = torch.tensor( - [emitted[:max_new_tokens]], + [session.output_ids], dtype=torch.long, - device=input_ids.device, + device=session.input_ids.device, ) - return torch.cat([input_ids, new_ids], dim=1) + return torch.cat([session.input_ids, new_ids], dim=1) @torch.no_grad() def _generate_batched( @@ -1510,255 +1959,119 @@ def _generate_batched( max_new_tokens: Union[int, Sequence[int]], stop_token_ids: Optional[List[int]], attention_mask: Optional[torch.Tensor], + sampling_contexts: Optional[tuple[SamplingContext, ...]] = None, + stop_token_ids_by_row: Optional[tuple[tuple[int, ...], ...]] = None, ) -> torch.Tensor: - """Track-C batched loop: one prefill + one verify per step for all rows. - - Greedy correctness rests on two facts: (1) the emitted stream of a - greedy verify step is always the target's argmax continuation given - the committed prefix -- accepted drafts matched the target by - definition and the bonus/correction IS the target's argmax -- so - drafts (and hence batching) can only change HOW MANY tokens a step - commits, never WHICH tokens are emitted; (2) a row's verify logits - depend only on its own cache row, block prefix, and RoPE positions, - which the left-pad ``attention_mask`` + per-row ``position_ids`` - plumbing reproduce exactly. - - Per-sequence rollback (C1): HF ``DynamicCache`` is dense -- one - ``cumulative_length`` per layer and slot-index-based causal/sliding - masks (``create_causal_mask`` reads ``q_offset`` from the cache, not - per-row positions) -- so rows cannot physically hold ragged lengths. - Instead every row rolls back to ``prev_start + min_cc`` where - ``min_cc`` is the SMALLEST per-row commit among still-active rows - (``rollback_target_cache`` is reused unchanged: the sliding-window - snapshot/rebuild is per-row inside the batched tensors). Rows that - committed more than ``min_cc`` carry the un-cached tail of their - already-emitted (hence target-true) tokens as the KNOWN PREFIX of - their next block (``build_block_with_prefixes``); the drafter only - fills the MASK slots past it, the verify re-caches the prefix, and - the prefix's re-confirmation tokens are skipped on emission - (``pending - 1`` of them). Slot distances equal true token distances - under this uniform-length scheme, so sliding-window masks stay exact. - """ - from transformers import DynamicCache - - batch, padded_prompt = int(input_ids.shape[0]), int(input_ids.shape[1]) - block_size = int(self.config.block_size) - layer_ids = list(self.config.target_layer_ids) - mask_token_id = int(self.config.mask_token_id) - - if isinstance(max_new_tokens, Sequence): - budgets = [int(x) for x in max_new_tokens] - if len(budgets) != batch: - raise ValueError( - f"per-sequence max_new_tokens has {len(budgets)} entries " - f"for batch size {batch}" - ) - else: - budgets = [int(max_new_tokens)] * batch - if any(b < 0 for b in budgets): - raise ValueError(f"max_new_tokens must be >= 0, got {budgets}") - - input_ids = input_ids.to(self.device) - self._configure_target_hooks(input_ids) + """Adapt the legacy tensor API to a driver-owned physical cohort.""" + from moe_infinity.spec_decode.backends import PhysicalCohortBackend + from moe_infinity.spec_decode.backends_bare_hf import ( + BatchedBareHFBackend, + ) + from moe_infinity.spec_decode.backends_rich import ( + BatchedRichBackend, + ) + from moe_infinity.spec_decode.session_driver import SessionDriver + if input_ids.ndim != 2: + raise ValueError( + "DFlashSpeculator._generate_batched expects input_ids of shape " + f"[batch, seq], got {tuple(input_ids.shape)}" + ) + batch = int(input_ids.shape[0]) + budgets = self._normalize_budgets(max_new_tokens, batch) + cohort_input_ids = input_ids.to(self.device) if attention_mask is None: - attention_mask = torch.ones_like(input_ids) + cohort_attention_mask = torch.ones_like(cohort_input_ids) else: - attention_mask = attention_mask.to( - device=self.device, dtype=torch.long - ) - if tuple(attention_mask.shape) != tuple(input_ids.shape): + cohort_attention_mask = attention_mask.to(device=self.device) + if tuple(cohort_attention_mask.shape) != tuple(cohort_input_ids.shape): raise ValueError( - f"attention_mask shape {tuple(attention_mask.shape)} != " - f"input_ids shape {tuple(input_ids.shape)}" + f"attention_mask shape {tuple(cohort_attention_mask.shape)} != " + f"input_ids shape {tuple(cohort_input_ids.shape)}" ) - if attention_mask.min() < 0 or attention_mask.max() > 1: + binary = (cohort_attention_mask == 0) | (cohort_attention_mask == 1) + if not bool(torch.all(binary).item()): raise ValueError("attention_mask must be 0/1 valued") - if int(attention_mask[:, -1].min()) != 1: - raise ValueError( - "batched DFlash requires LEFT-padded prompts: every row's last " - "token must be real (attention_mask[:, -1] == 1)" - ) - steps = attention_mask[:, 1:] - attention_mask[:, :-1] - if int(steps.min()) < 0: - raise ValueError( - "batched DFlash requires LEFT-padded prompts: each " - "attention_mask row must be 0*1* (pads first, then real tokens)" - ) - pads = padded_prompt - attention_mask.sum(dim=1) + cohort_attention_mask = cohort_attention_mask.to(dtype=torch.long) - prefill_position_ids = (attention_mask.cumsum(dim=-1) - 1).clamp_min(0) - logits, hidden_states, target_kv = self._forward_target( - input_ids, - past_key_values=None, - logits_to_keep=1, - attention_mask=attention_mask, - position_ids=prefill_position_ids, - ) - anchors = logits[:, -1, :].argmax(dim=-1) - context_feature = _extract_context_feature(hidden_states, layer_ids).to( - self.device - ) - stop_ids = set(_resolve_stop_ids(self.target, stop_token_ids)) - - emitted: List[List[int]] = [[int(anchors[b])] for b in range(batch)] - finished: List[bool] = [ - budgets[b] <= 0 or (bool(stop_ids) and int(anchors[b]) in stop_ids) - for b in range(batch) - ] - start = padded_prompt - draft_kv = DynamicCache() if self._drafter_has_kv_cache else None - self.step_trace = [] - - def active(b: int) -> bool: - return not finished[b] and len(emitted[b]) < budgets[b] + if sampling_contexts is None: + normalized_sampling = tuple(SamplingContext() for _ in range(batch)) + else: + normalized_sampling = tuple(sampling_contexts) + if len(normalized_sampling) != batch: + raise ValueError( + f"sampling_contexts has {len(normalized_sampling)} entries " + f"for batch size {batch}" + ) + if stop_token_ids_by_row is None: + shared_stops = tuple(_resolve_stop_ids(self.target, stop_token_ids)) + stop_rows = tuple(shared_stops for _ in range(batch)) + else: + stop_rows = tuple(tuple(row) for row in stop_token_ids_by_row) + if len(stop_rows) != batch: + raise ValueError( + f"stop_token_ids_by_row has {len(stop_rows)} entries " + f"for batch size {batch}" + ) - while any(active(b) for b in range(batch)): - prev_start = start - pendings = [ - len(emitted[b]) - (start - padded_prompt) for b in range(batch) - ] - prefixes = [ - emitted[b][start - padded_prompt :] if active(b) else [] - for b in range(batch) + requests: list[RequestSpec] = [] + for row in range(batch): + row_ids = cohort_input_ids[row][ + cohort_attention_mask[row].to(dtype=torch.bool) ] - block = build_block_with_prefixes( - prefixes, mask_token_id, block_size - ).to(self.device) - - drafter_out = self._run_drafter( - block, context_feature, start, draft_kv - ) - draft_logits = self.lm_head(drafter_out)[:, -(block_size - 1) :, :] - for b in range(batch): - if active(b) and pendings[b] < block_size: - block[b, pendings[b] :] = draft_logits[ - b, pendings[b] - 1 : - ].argmax(dim=-1) - - cache_snapshot = snapshot_target_cache(target_kv) - - block_attention = torch.cat( - [ - attention_mask, - torch.ones( - batch, - start - padded_prompt + block_size, - dtype=attention_mask.dtype, - device=self.device, + requests.append( + RequestSpec( + request_id=f"direct-{row}", + prompt_token_ids=tuple( + int(token) for token in row_ids.tolist() ), - ], - dim=1, - ) - block_position_ids = torch.arange( - start, start + block_size, device=self.device, dtype=torch.long - ).unsqueeze(0) - pads.unsqueeze(1) - logits, hidden_states, target_kv = self._verify_target_block( - block, - target_kv, - attention_mask=block_attention, - position_ids=block_position_ids, - ) - posterior = logits.argmax(dim=-1).to(self.device) - accepts = acceptance_lengths(block, posterior) - step_committed = committed_tokens_ragged(block, posterior, accepts) - - # Per-row emission with the re-fed prefix skipped: of this step's - # ``accept + 1`` emitted tokens the first ``pending - 1`` are - # re-confirmations of tokens already emitted (the known prefix), - # so row b newly emits ``step_tokens[pending - 1:]``. Keeping k of - # those commits ``min(pending - 1 + k, accept) + 1`` cached tokens - # (the v1 rule ``min(k, accept) + 1`` at pending == 1). - step_cc: dict[int, int] = {} - for b in range(batch): - if not active(b): - continue - pending = pendings[b] - accept = accepts[b] - step_tokens = [ - int(t) for t in step_committed[b].emitted[0].tolist() - ] - new_tokens = step_tokens[pending - 1 :] - keep = len(new_tokens) - stop = False - if stop_ids: - for j, tok in enumerate(new_tokens): - if tok in stop_ids: - keep = j + 1 - stop = True - break - remaining = budgets[b] - len(emitted[b]) - if keep > remaining: - keep = remaining - stop = True - emitted[b].extend(new_tokens[:keep]) - step_cc[b] = min(pending - 1 + keep, accept) + 1 - if stop or len(emitted[b]) >= budgets[b]: - finished[b] = True - - continuing = [b for b in step_cc if active(b)] - min_cc = ( - min(step_cc[b] for b in continuing) - if continuing - else min(step_cc.values()) - ) - rollback_target_cache( - target_kv, - cache_snapshot, - prev_start=prev_start, - committed=min_cc, - block_size=block_size, - ) - start = prev_start + min_cc - assert int(target_kv.get_seq_length()) == start - - if self.route_ahead_stats is not None: - self.route_ahead_stats.commit_step(kept_rows=min_cc) - - for b, cc_b in step_cc.items(): - self.step_trace.append( - NativeStepTrace( - prev_start=prev_start, - accept=cc_b - 1, - start=start, - emitted_len=len(emitted[b]), - target_cache_len=int(target_kv.get_seq_length()), - draft_cache_len=( - int(draft_kv.get_seq_length()) - if draft_kv is not None - else None - ), - ) + max_new_tokens=budgets[row], + stop_token_ids=frozenset(stop_rows[row]), + sampling=normalized_sampling[row], ) - if not continuing: - break - - suffix = _extract_context_feature(hidden_states, layer_ids).to( - self.device - )[:, :min_cc, :] - if self._drafter_has_kv_cache: - context_feature = suffix - else: - context_feature = torch.cat([context_feature, suffix], dim=1) - - self.last_target_cache = target_kv - self.last_draft_cache = draft_kv + ) - new_lengths = [min(len(emitted[b]), budgets[b]) for b in range(batch)] + backend = ( + BatchedRichBackend(self) + if callable(getattr(self.moe, "_native_model_forward_rich", None)) + else BatchedBareHFBackend(self) + ) + run = SessionDriver( + [cast(PhysicalCohortBackend, backend)] + ).run_physical_cohort( + cohort_input_ids, + requests=tuple(requests), + attention_mask=cohort_attention_mask, + ) + result = run.backend_result + + self.step_trace = list(result.step_trace) + self.last_target_cache = result.target_cache + self.last_draft_cache = result.draft_cache + self.last_session_results = run.results + self.last_session_traces = tuple(row.trace for row in run.results) + new_lengths = [len(row.output_token_ids) for row in run.results] self.last_generated_lengths = new_lengths pad_id = _get(_get(self.target, "config", self.target), "pad_token_id") pad_id = 0 if pad_id is None else int(pad_id) width = max(new_lengths) if new_lengths else 0 new_ids = torch.full( - (batch, width), pad_id, dtype=torch.long, device=input_ids.device + (batch, width), + pad_id, + dtype=cohort_input_ids.dtype, + device=cohort_input_ids.device, ) - for b in range(batch): - n = new_lengths[b] - if n: - new_ids[b, :n] = torch.tensor( - emitted[b][:n], dtype=torch.long, device=input_ids.device + for row, driver_result in enumerate(run.results): + generated = driver_result.output_token_ids + if generated: + new_ids[row, : len(generated)] = torch.tensor( + generated, + dtype=cohort_input_ids.dtype, + device=cohort_input_ids.device, ) - return torch.cat([input_ids, new_ids], dim=1) + return torch.cat([cohort_input_ids, new_ids], dim=1).to( + input_ids.device + ) def run( self, @@ -1783,7 +2096,6 @@ def run( generated token ids (prompt stripped); the engine wraps them in a ``GenerationResult`` and ``MoE.generate`` re-prepends the prompt. """ - del request_id # protocol-conformance parameter; the loop is stateless input_ids = torch.tensor([list(prompt_token_ids)], dtype=torch.long) max_new_tokens = int(getattr(sampling_params, "max_tokens", 256)) eos = getattr(engine, "eos_token_id", None) @@ -1798,6 +2110,29 @@ def run( ) return [int(t) for t in output[0, len(prompt_token_ids) :].tolist()] + def supports_engine_request( + self, sampling_params: SamplingParams, *, batch_size: int + ) -> bool: + """Return whether the compatibility engine may select DFlash.""" + if batch_size != 1: + return False + sampling = SamplingContext( + temperature=float(getattr(sampling_params, "temperature", 0.0)), + top_k=int(getattr(sampling_params, "top_k", 0) or 0), + top_p=float(getattr(sampling_params, "top_p", 1.0)), + ) + if sampling.is_sampled or getattr(sampling_params, "do_sample", False): + return False + from moe_infinity.spec_decode.backends import DFlashExecutionBackend + + request = RequestSpec( + request_id="engine-capability-check", + prompt_token_ids=(0,), + max_new_tokens=int(getattr(sampling_params, "max_tokens", 256)), + sampling=sampling, + ) + return DFlashExecutionBackend(self).supports(request) + __all__ = [ "DFlashConfig", @@ -1807,6 +2142,8 @@ def run( "SpecSession", "VerifyResult", "bind_shared_weights", + "build_pairing_evidence", + "executor_wiring_reachable", "project_expert_bytes", "read_dflash_config", "validate_drafter", diff --git a/moe_infinity/spec_decode/protocols.py b/moe_infinity/spec_decode/protocols.py new file mode 100644 index 00000000..7780acc2 --- /dev/null +++ b/moe_infinity/spec_decode/protocols.py @@ -0,0 +1,586 @@ +"""Shared contracts for DFlash execution backends and request sessions. + +The types in this module are deliberately model- and scheduler-independent. +They define the values exchanged by native and serving implementations without +requiring either implementation to own a second sampling, cache, or trace +schema. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Any, Literal, NamedTuple, Protocol, runtime_checkable + +import torch + +CacheKind = Literal["dense_dynamic", "paged", "other"] + + +@dataclass(frozen=True) +class PairingEvidence: + """Structural target/drafter compatibility, independent of execution. + + The fields are limited to the authoritative DFlash config, shape, vocab, + mask, layer, block, and drafter-module checks. Executor wiring and + route-ahead observations deliberately live in :class:`ExecutorEvidence`. + ``validated_checkpoint_scope`` is empty unless the caller explicitly knows + the checkpoint identities covered by the validation. + """ + + valid: bool = False + config_valid: bool = False + dimensions_valid: bool = False + vocab_valid: bool = False + mask_valid: bool = False + layers_valid: bool = False + block_valid: bool = False + module_valid: bool | None = None + validated_checkpoint_scope: tuple[str, ...] = () + failure_reason: str | None = None + + def __post_init__(self) -> None: + if self.valid and not all( + ( + self.config_valid, + self.dimensions_valid, + self.vocab_valid, + self.mask_valid, + self.layers_valid, + self.block_valid, + self.module_valid is not False, + ) + ): + raise ValueError( + "valid pairing evidence requires all checked fields" + ) + if any(not item for item in self.validated_checkpoint_scope): + raise ValueError( + "validated_checkpoint_scope entries must be non-empty" + ) + + def as_dict(self) -> dict[str, object]: + result: dict[str, object] = { + "valid": self.valid, + "config_valid": self.config_valid, + "dimensions_valid": self.dimensions_valid, + "vocab_valid": self.vocab_valid, + "mask_valid": self.mask_valid, + "layers_valid": self.layers_valid, + "block_valid": self.block_valid, + "module_valid": self.module_valid, + "validated_checkpoint_scope": self.validated_checkpoint_scope, + "failure_reason": self.failure_reason, + } + return result + + +@dataclass(frozen=True) +class ExecutorEvidence: + """Observed executor reachability and route-ahead behavior. + + This is an immutable snapshot. It never establishes target/drafter + compatibility and may therefore report a reachable executor for an invalid + pair, or an unreachable executor for a valid published pair. + """ + + wiring_reachable: bool = False + prefetcher_present: bool = False + attempted_layers: tuple[int, ...] = () + fired_layers: tuple[int, ...] = () + actual_expert_union: frozenset[tuple[int, int]] = frozenset() + actual_expert_union_by_row: frozenset[tuple[int, int, int]] = frozenset() + prefetched_bytes: int = 0 + coverage: float | None = None + wasted_prefetch_bytes: int | None = None + cache_hit_rate: float | None = None + fallback_reason: str | None = None + + def __post_init__(self) -> None: + if any( + layer < 0 for layer in self.attempted_layers + self.fired_layers + ): + raise ValueError("executor evidence layer ids must be non-negative") + if not set(self.fired_layers).issubset(self.attempted_layers): + raise ValueError( + "fired_layers must be a subset of attempted_layers" + ) + if any( + layer < 0 or expert < 0 + for layer, expert in self.actual_expert_union + ): + raise ValueError("actual expert ids must be non-negative") + if any( + row < 0 or layer < 0 or expert < 0 + for row, layer, expert in self.actual_expert_union_by_row + ): + raise ValueError("row-aware actual expert ids must be non-negative") + if self.prefetched_bytes < 0: + raise ValueError("prefetched_bytes must be >= 0") + if ( + self.wasted_prefetch_bytes is not None + and self.wasted_prefetch_bytes < 0 + ): + raise ValueError("wasted_prefetch_bytes must be >= 0") + for name, value in ( + ("coverage", self.coverage), + ("cache_hit_rate", self.cache_hit_rate), + ): + if value is not None and not 0.0 <= value <= 1.0: + raise ValueError(f"{name} must be in [0, 1]") + + def as_dict(self) -> dict[str, object]: + result: dict[str, object] = { + "wiring_reachable": self.wiring_reachable, + "prefetcher_present": self.prefetcher_present, + "attempted_layers": self.attempted_layers, + "fired_layers": self.fired_layers, + "actual_expert_union": tuple(sorted(self.actual_expert_union)), + "prefetched_bytes": self.prefetched_bytes, + "coverage": self.coverage, + "wasted_prefetch_bytes": self.wasted_prefetch_bytes, + "cache_hit_rate": self.cache_hit_rate, + "fallback_reason": self.fallback_reason, + } + if self.actual_expert_union_by_row: + result["actual_expert_union_by_row"] = tuple( + sorted(self.actual_expert_union_by_row) + ) + return result + + +@dataclass(frozen=True) +class SamplingContext: + """Request-scoped sampling policy and random-number stream.""" + + temperature: float = 0.0 + top_k: int = 0 + top_p: float = 1.0 + generator: torch.Generator | None = None + + def __post_init__(self) -> None: + if self.temperature < 0: + raise ValueError("temperature must be >= 0") + if self.top_k < 0: + raise ValueError("top_k must be >= 0") + if not 0 < self.top_p <= 1: + raise ValueError("top_p must be in (0, 1]") + + @property + def is_greedy(self) -> bool: + return self.temperature == 0 + + @property + def is_sampled(self) -> bool: + return not self.is_greedy + + +@dataclass(frozen=True) +class RequestSpec: + """Backend-neutral inputs needed to start one speculative session.""" + + request_id: str + prompt_token_ids: tuple[int, ...] + max_new_tokens: int + stop_token_ids: frozenset[int] = frozenset() + sampling: SamplingContext = field(default_factory=SamplingContext) + + def __post_init__(self) -> None: + if not self.request_id: + raise ValueError("request_id must be non-empty") + if not self.prompt_token_ids: + raise ValueError("prompt_token_ids must be non-empty") + if self.max_new_tokens < 0: + raise ValueError("max_new_tokens must be >= 0") + + @property + def prompt_length(self) -> int: + return len(self.prompt_token_ids) + + @property + def is_sampled(self) -> bool: + return self.sampling.is_sampled + + +class NativeStepTrace(NamedTuple): + """Canonical per-round cache and emission accounting. + + ``accept`` counts accepted draft tokens. The anchor plus those drafts are + cached, while the emitted bonus remains uncached for the next round. + """ + + prev_start: int + accept: int + start: int + emitted_len: int + target_cache_len: int + draft_cache_len: int | None + + @property + def committed_count(self) -> int: + """Number of block tokens retained in cache by this round.""" + + return self.accept + 1 + + @property + def cache_advance(self) -> int: + return self.start - self.prev_start + + +@dataclass(frozen=True) +class SessionRoundResult: + """Backend-neutral result of one draft/verify/commit round. + + ``accepted_draft_count`` records the verify rule's untruncated acceptance. + A finished stop/budget-truncated round may emit fewer tokens, including no + tokens for a defensive no-op entry. + """ + + accepted_draft_count: int + committed_token_ids: tuple[int, ...] + next_anchor: int | None + target_cache_length: int + emitted_length: int + finished: bool + finish_reason: str | None + fallback_reason: str | None = None + + def __post_init__(self) -> None: + if self.accepted_draft_count < 0: + raise ValueError("accepted_draft_count must be >= 0") + if self.target_cache_length < 0: + raise ValueError("target_cache_length must be >= 0") + if self.emitted_length < 0: + raise ValueError("emitted_length must be >= 0") + committed_count = len(self.committed_token_ids) + complete_count = self.accepted_draft_count + 1 + if committed_count == 0: + if not self.finished or self.next_anchor is not None: + raise ValueError( + "an empty committed_token_ids result must be a finished no-op " + "with no next_anchor" + ) + elif committed_count != complete_count and ( + not self.finished or committed_count > complete_count + ): + raise ValueError( + "committed_token_ids must contain at most the accepted drafts " + "plus one target bonus token; only finished rounds may be " + "stop/budget-truncated" + ) + + @property + def cached_token_count(self) -> int: + """Anchor plus accepted drafts retained by the target cache.""" + + if not self.committed_token_ids: + return 0 + return min(len(self.committed_token_ids), self.accepted_draft_count) + 1 + + @property + def emitted_token_count(self) -> int: + return len(self.committed_token_ids) + + @property + def commit_block_token_ids(self) -> tuple[int, ...]: + """Tokens emitted by this commit block after stop/budget truncation.""" + + return self.committed_token_ids + + @property + def accepted_drafts(self) -> int: + """Compatibility alias for the verify rule's accepted-draft count.""" + + return self.accepted_draft_count + + @property + def bonus_token_id(self) -> int | None: + return self.next_anchor + + +@dataclass(frozen=True) +class BackendCapabilities: + """Static execution features advertised by a speculative backend.""" + + supports_batch: bool + supports_sampling: bool + supports_ragged_rows: bool + cache_kind: CacheKind + supports_route_ahead: bool + supports_rich_forward: bool + pairing_evidence: PairingEvidence = field(default_factory=PairingEvidence) + executor_evidence: ExecutorEvidence = field( + default_factory=ExecutorEvidence + ) + + def __post_init__(self) -> None: + if self.cache_kind not in ("dense_dynamic", "paged", "other"): + raise ValueError( + "cache_kind must be 'dense_dynamic', 'paged', or 'other'" + ) + + @property + def is_dense_cache(self) -> bool: + return self.cache_kind == "dense_dynamic" + + @property + def is_paged_cache(self) -> bool: + return self.cache_kind == "paged" + + +@dataclass(frozen=True) +class CacheSnapshot: + """Rollback point shared by cache adapters.""" + + logical_length: int + + def __post_init__(self) -> None: + if self.logical_length < 0: + raise ValueError("logical_length must be >= 0") + + +@dataclass(frozen=True) +class RichBatchMetadata: + """Row layout and model inputs for one physical rich forward.""" + + row_offsets: tuple[int, ...] + row_lengths: tuple[int, ...] + attention_mask: torch.Tensor | None = None + position_ids: torch.Tensor | None = None + cache_handles: tuple[object, ...] = () + request_contexts: tuple[object, ...] = () + route_contexts: tuple[object, ...] = () + block_tables: torch.Tensor | None = None + slot_mapping: torch.Tensor | None = None + seq_lens: torch.Tensor | None = None + is_prefill: bool = True + + def __post_init__(self) -> None: + rows = len(self.row_lengths) + if len(self.row_offsets) != rows + 1 or self.row_offsets[:1] != (0,): + raise ValueError( + "row_offsets must start at zero and have one sentinel" + ) + expected = 0 + for row, length in enumerate(self.row_lengths): + if length < 0: + raise ValueError("row_lengths must be non-negative") + expected += length + if self.row_offsets[row + 1] != expected: + raise ValueError("row_offsets must be cumulative row_lengths") + for name, values in ( + ("cache_handles", self.cache_handles), + ("request_contexts", self.request_contexts), + ("route_contexts", self.route_contexts), + ): + if values and len(values) != rows: + raise ValueError(f"{name} must have one entry per row") + for name, value in ( + ("attention_mask", self.attention_mask), + ("position_ids", self.position_ids), + ): + if value is not None and ( + value.ndim != 2 or value.shape[0] != rows + ): + raise ValueError(f"{name} must have shape [rows, sequence]") + if self.block_tables is not None and self.block_tables.shape[0] != rows: + raise ValueError("block_tables must have one row per request") + if self.seq_lens is not None and self.seq_lens.numel() != rows: + raise ValueError("seq_lens must have one entry per row") + + @property + def row_count(self) -> int: + return len(self.row_lengths) + + +@dataclass(frozen=True) +class RichForwardResult: + """Target forward payload whose cache handle may be engine-owned.""" + + logits: torch.Tensor + hidden_states: tuple[torch.Tensor, ...] + cache_handle: object + cache_handles: tuple[object, ...] = () + row_offsets: tuple[int, ...] = () + row_lengths: tuple[int, ...] = () + + def __post_init__(self) -> None: + if not self.cache_handles: + object.__setattr__(self, "cache_handles", (self.cache_handle,)) + if self.row_lengths: + if len(self.row_offsets) != len(self.row_lengths) + 1: + raise ValueError( + "row_offsets must have one sentinel per rich row" + ) + expected = 0 + for row, length in enumerate(self.row_lengths): + expected += length + if self.row_offsets[row + 1] != expected: + raise ValueError( + "row_offsets must be cumulative row_lengths" + ) + if len(self.cache_handles) not in (1, len(self.row_lengths)): + raise ValueError("cache_handles must be shared or row-aligned") + + def __iter__(self): + yield self.logits + yield self.hidden_states + yield self.cache_handle + + +@runtime_checkable +class CacheAdapter(Protocol): + """Structural cache lifecycle required by speculative sessions.""" + + def snapshot(self) -> CacheSnapshot: + """Capture the cache position before a tentative verify.""" + ... + + def restore(self, snapshot: CacheSnapshot) -> None: + """Restore an earlier logical cache position.""" + ... + + def append(self, token_count: int) -> None: + """Record tokens appended by a completed model forward.""" + ... + + def truncate(self, logical_length: int) -> None: + """Discard cache state beyond ``logical_length``.""" + ... + + def logical_length(self) -> int: + """Return the current logical token count.""" + ... + + def release(self) -> None: + """Release all cache state owned by this adapter.""" + ... + + +@runtime_checkable +class _DenseCache(Protocol): + def get_seq_length(self) -> int: ... + + def crop(self, length: int) -> None: ... + + +class DenseCacheAdapter: + """Adapter for transformers-style dense caches with ``crop`` support.""" + + cache_kind: CacheKind = "dense_dynamic" + + def __init__(self, cache: object) -> None: + if not isinstance(cache, _DenseCache): + raise TypeError( + "dense cache must provide get_seq_length() and crop()" + ) + self.cache: _DenseCache = cache + self._logical_length: int = int(cache.get_seq_length()) + self._released: bool = False + + def _ensure_active(self) -> None: + if self._released: + raise RuntimeError("dense cache adapter has been released") + + def snapshot(self) -> CacheSnapshot: + self._ensure_active() + return CacheSnapshot(logical_length=self._logical_length) + + def restore(self, snapshot: CacheSnapshot) -> None: + self.truncate(snapshot.logical_length) + + def append(self, token_count: int) -> None: + self._ensure_active() + if token_count < 0: + raise ValueError("token_count must be >= 0") + expected_length = self._logical_length + token_count + physical_length = int(self.cache.get_seq_length()) + if physical_length != expected_length: + raise RuntimeError( + "dense cache physical length does not match the reported append: " + + f"expected {expected_length}, got {physical_length}" + ) + self._logical_length = expected_length + + def truncate(self, logical_length: int) -> None: + self._ensure_active() + if logical_length < 0 or logical_length > self._logical_length: + raise ValueError( + "logical_length must be between 0 and the current logical length" + ) + self.cache.crop(logical_length) + self._logical_length = logical_length + + def logical_length(self) -> int: + self._ensure_active() + return self._logical_length + + def release(self) -> None: + if self._released: + return + self.cache.crop(0) + self._logical_length = 0 + self._released = True + + +@dataclass +class SessionTrace: + """Serializable session accounting built from ``NativeStepTrace`` rows.""" + + request_id: str + backend: str + cache_kind: CacheKind + sampled: bool + round_count: int = 0 + accepted: int = 0 + committed: int = 0 + emitted: int = 0 + rollback: int = 0 + replay: int = 0 + finish_reason: str | None = None + route_ahead_status: str | None = None + pairing_evidence: PairingEvidence = field(default_factory=PairingEvidence) + executor_evidence: ExecutorEvidence = field( + default_factory=ExecutorEvidence + ) + + def append(self, step: NativeStepTrace) -> None: + self.round_count += 1 + self.accepted += step.accept + self.committed += step.committed_count + self.emitted = step.emitted_len + + def as_dict(self) -> dict[str, object]: + return { + "request_id": self.request_id, + "backend": self.backend, + "cache_kind": self.cache_kind, + "sampled": self.sampled, + "round_count": self.round_count, + "accepted": self.accepted, + "committed": self.committed, + "emitted": self.emitted, + "rollback": self.rollback, + "replay": self.replay, + "finish_reason": self.finish_reason, + "route_ahead_status": self.route_ahead_status, + "pairing_evidence": self.pairing_evidence.as_dict(), + "executor_evidence": self.executor_evidence.as_dict(), + } + + +__all__ = [ + "BackendCapabilities", + "CacheAdapter", + "CacheKind", + "CacheSnapshot", + "DenseCacheAdapter", + "ExecutorEvidence", + "NativeStepTrace", + "PairingEvidence", + "RichBatchMetadata", + "RichForwardResult", + "RequestSpec", + "SamplingContext", + "SessionRoundResult", + "SessionTrace", +] diff --git a/moe_infinity/spec_decode/session_driver.py b/moe_infinity/spec_decode/session_driver.py new file mode 100644 index 00000000..03fa86e1 --- /dev/null +++ b/moe_infinity/spec_decode/session_driver.py @@ -0,0 +1,599 @@ +"""Deterministic request cohort planning and singleton session execution.""" + +from __future__ import annotations + +from collections import OrderedDict +from collections.abc import Iterable, Sequence +from dataclasses import dataclass +from types import TracebackType +from typing import Any, Hashable, cast + +import torch + +from .backends import ( + ExecutionBackend, + PhysicalCohortBackend, + PhysicalCohortResult, +) +from .protocols import ( + BackendCapabilities, + RequestSpec, + SamplingContext, + SessionRoundResult, + SessionTrace, +) + + +class UnsupportedRequestError(RuntimeError): + """No configured backend can preserve a request's semantics.""" + + +class BackendProgressError(RuntimeError): + """A backend returned an unfinished round without observable progress.""" + + +class SessionCleanupError(RuntimeError): + """One or more cleanup operations failed after execution completed.""" + + def __init__(self, errors: Sequence[BaseException]) -> None: + self.errors = tuple(errors) + super().__init__( + f"{len(self.errors)} session cleanup operation(s) failed: " + + "; ".join(str(error) for error in self.errors) + ) + + +def _attach_cleanup_context( + primary: BaseException, errors: Sequence[BaseException] +) -> None: + """Retain cleanup failures without replacing ``primary``. + + Python 3.11+ exceptions normally expose ``add_note``. Python 3.10 and + custom exceptions may not, so the fallback chains an aggregate cleanup + error through ``__context__`` and exposes the original errors as an + attribute when possible. Cleanup reporting itself must never mask the + execution failure. + """ + if not errors: + return + + try: + add_note = getattr(primary, "add_note", None) + except BaseException: + add_note = None + if callable(add_note): + try: + for error in errors: + add_note(f"session cleanup failed: {error}") + return + except BaseException: + pass + + cleanup_context = SessionCleanupError(errors) + try: + primary.__context__ = cleanup_context + except BaseException: + pass + try: + setattr(primary, "session_cleanup_errors", tuple(errors)) + except BaseException: + pass + + +@dataclass(frozen=True) +class CohortPlan: + """Stable row indices sharing one selected backend compatibility key.""" + + backend: str + capabilities: BackendCapabilities + compatibility_key: Hashable + row_indices: tuple[int, ...] + + +@dataclass(frozen=True) +class DriverResult: + """One completed request result; returned in original row order.""" + + request_id: str + output_token_ids: tuple[int, ...] + finish_reason: str + backend: str + trace: SessionTrace + rounds: tuple[SessionRoundResult, ...] + fallback_reason: str | None = None + + +@dataclass(frozen=True) +class PhysicalCohortDriverResult: + """Atomically published row results plus backend-owned diagnostics.""" + + results: tuple[DriverResult, ...] + backend_result: PhysicalCohortResult + + +@dataclass(frozen=True) +class _Selection: + row_index: int + request: RequestSpec + backend: ExecutionBackend[Any, Any] + compatibility_key: Hashable + fallback_reason: str | None + + +@dataclass(frozen=True) +class _PhysicalSelection: + row_index: int + request: RequestSpec + backend: PhysicalCohortBackend + compatibility_key: Hashable + fallback_reason: str | None + + +@dataclass +class _ActiveSession: + selection: _Selection + session: object + initial_snapshot: object + rounds: list[SessionRoundResult] + observable_progress: int + + +class SessionDriver: + """Plan compatible cohorts, then drive one backend session per request. + + Backend order is the documented fallback policy: the first backend that + both advertises the required capability and accepts the request wins. + Sampling requests skip non-sampling backends; they are never rewritten as + greedy requests. Selection for every row completes before the first + prefill, making unsupported input a pre-output failure. + """ + + def __init__( + self, + backends: Sequence[ExecutionBackend[Any, Any] | PhysicalCohortBackend], + ) -> None: + if not backends: + raise ValueError("SessionDriver requires at least one backend") + self.backends = tuple(backends) + self.last_cohorts: tuple[CohortPlan, ...] = () + self.last_results: tuple[DriverResult, ...] = () + + def run( + self, requests: RequestSpec | Iterable[RequestSpec] + ) -> tuple[DriverResult, ...]: + normalized = self._normalize_requests(requests) + selections = tuple( + self._select_backend(row_index, request) + for row_index, request in enumerate(normalized) + ) + cohorts = self._cohort(selections) + self.last_cohorts = tuple( + CohortPlan( + backend=rows[0].backend.name, + capabilities=rows[0].backend.capabilities, + compatibility_key=rows[0].compatibility_key, + row_indices=tuple(row.row_index for row in rows), + ) + for rows in cohorts + ) + self.last_results = () + + active: list[_ActiveSession] = [] + pending_results: tuple[DriverResult, ...] | None = None + primary: BaseException | None = None + primary_traceback: TracebackType | None = None + try: + # Prefill remains singleton, but happens only after the complete + # capability plan has succeeded for every row. + for selection in selections: + session = selection.backend.prefill(selection.request) + initial_progress = len( + selection.backend.output_token_ids(session) + ) + active.append( + _ActiveSession( + selection=selection, + session=session, + initial_snapshot=selection.backend.snapshot(session), + rounds=[], + observable_progress=initial_progress, + ) + ) + + by_row = {item.selection.row_index: item for item in active} + for cohort in cohorts: + rows = [by_row[selection.row_index] for selection in cohort] + self._run_cohort(rows) + + pending_results = tuple( + self._result(item) + for item in sorted( + active, key=lambda item: item.selection.row_index + ) + ) + except BaseException as exc: + primary = exc + primary_traceback = exc.__traceback__ + + cleanup_errors = self._cleanup(active, restore=primary is not None) + if primary is not None: + _attach_cleanup_context(primary, cleanup_errors) + raise primary.with_traceback(primary_traceback) + if cleanup_errors: + raise SessionCleanupError(cleanup_errors) + + assert pending_results is not None + self.last_results = pending_results + return pending_results + + def run_physical_cohort( + self, + input_ids: torch.Tensor, + *, + requests: RequestSpec | Iterable[RequestSpec], + attention_mask: torch.Tensor, + ) -> PhysicalCohortDriverResult: + """Select and execute one physical backend cohort atomically. + + All rows are capability-selected before the backend receives the + cohort. A physical entry never splits rows or rewrites sampling policy; + callers that need semantic per-request fallback must use :meth:`run`. + """ + self.last_results = () + self.last_cohorts = () + normalized = self._normalize_requests(requests) + self._validate_physical_inputs( + input_ids, attention_mask, len(normalized) + ) + if not normalized: + raise ValueError("a physical cohort requires at least one request") + + selections = tuple( + self._select_physical_backend(row_index, request) + for row_index, request in enumerate(normalized) + ) + first = selections[0] + if any( + selection.backend is not first.backend + or selection.compatibility_key != first.compatibility_key + for selection in selections[1:] + ): + raise UnsupportedRequestError( + "requests cannot execute as one physical cohort without " + "splitting backends or compatibility keys" + ) + + self.last_cohorts = ( + CohortPlan( + backend=first.backend.name, + capabilities=first.backend.capabilities, + compatibility_key=first.compatibility_key, + row_indices=tuple(range(len(normalized))), + ), + ) + stop_rows = tuple( + tuple(sorted(request.stop_token_ids)) for request in normalized + ) + shared_stops = stop_rows[0] + common_stops = all(row == shared_stops for row in stop_rows) + budgets = tuple(request.max_new_tokens for request in normalized) + sampling_contexts = tuple(request.sampling for request in normalized) + sampled = any(context.is_sampled for context in sampling_contexts) + backend_result = self._execute_physical_backend( + first.backend, + input_ids, + budgets=budgets, + shared_stops=shared_stops if common_stops else (), + attention_mask=attention_mask, + sampling_contexts=sampling_contexts if sampled else None, + stop_rows=None if common_stops else stop_rows, + ) + results = self._physical_results(selections, backend_result) + self.last_results = results + return PhysicalCohortDriverResult( + results=results, + backend_result=backend_result, + ) + + @staticmethod + def _normalize_requests( + requests: RequestSpec | Iterable[RequestSpec], + ) -> tuple[RequestSpec, ...]: + rows = ( + (requests,) + if isinstance(requests, RequestSpec) + else tuple(requests) + ) + if not rows: + return () + if any(not isinstance(row, RequestSpec) for row in rows): + raise TypeError("SessionDriver.run expects RequestSpec rows") + request_ids = [row.request_id for row in rows] + if len(set(request_ids)) != len(request_ids): + raise ValueError( + "request_id values must be unique within a driver run" + ) + return rows + + def _select_backend( + self, row_index: int, request: RequestSpec + ) -> _Selection: + for backend_index, candidate in enumerate(self.backends): + if not isinstance(candidate, ExecutionBackend): + continue + backend = cast(ExecutionBackend[Any, Any], candidate) + if ( + request.is_sampled + and not backend.capabilities.supports_sampling + ): + continue + if not backend.supports(request): + continue + compatibility_key = backend.cohort_key(request) + try: + hash(compatibility_key) + except TypeError as exc: + raise TypeError("backend cohort_key must be hashable") from exc + fallback_reason = ( + None + if backend_index == 0 + else f"selected compatible fallback backend {backend.name}" + ) + return _Selection( + row_index=row_index, + request=request, + backend=backend, + compatibility_key=compatibility_key, + fallback_reason=fallback_reason, + ) + mode = "sampled" if request.is_sampled else "greedy" + raise UnsupportedRequestError( + f"request {request.request_id!r} has no compatible {mode} backend" + ) + + def _select_physical_backend( + self, row_index: int, request: RequestSpec + ) -> _PhysicalSelection: + physical_index = 0 + for candidate in self.backends: + if not isinstance(candidate, PhysicalCohortBackend): + continue + backend = cast(PhysicalCohortBackend, candidate) + if not backend.capabilities.supports_batch: + physical_index += 1 + continue + if ( + request.is_sampled + and not backend.capabilities.supports_sampling + ): + physical_index += 1 + continue + if not backend.supports(request): + physical_index += 1 + continue + compatibility_key = backend.cohort_key(request) + try: + hash(compatibility_key) + except TypeError as exc: + raise TypeError("backend cohort_key must be hashable") from exc + return _PhysicalSelection( + row_index=row_index, + request=request, + backend=backend, + compatibility_key=compatibility_key, + fallback_reason=( + None + if physical_index == 0 + else f"selected compatible fallback backend {backend.name}" + ), + ) + mode = "sampled" if request.is_sampled else "greedy" + raise UnsupportedRequestError( + f"request {request.request_id!r} has no compatible physical {mode} backend" + ) + + @staticmethod + def _validate_physical_inputs( + input_ids: torch.Tensor, + attention_mask: torch.Tensor, + request_count: int, + ) -> None: + if input_ids.ndim != 2: + raise ValueError( + "physical cohort input_ids must have shape [batch, seq]" + ) + if tuple(attention_mask.shape) != tuple(input_ids.shape): + raise ValueError( + f"attention_mask shape {tuple(attention_mask.shape)} != " + f"input_ids shape {tuple(input_ids.shape)}" + ) + if int(input_ids.shape[0]) != request_count: + raise ValueError( + f"physical cohort has {int(input_ids.shape[0])} tensor rows " + f"for {request_count} requests" + ) + binary = (attention_mask == 0) | (attention_mask == 1) + if not bool(torch.all(binary).item()): + raise ValueError("attention_mask must be 0/1 valued") + + @staticmethod + def _execute_physical_backend( + backend: PhysicalCohortBackend, + input_ids: torch.Tensor, + *, + budgets: tuple[int, ...], + shared_stops: tuple[int, ...], + attention_mask: torch.Tensor, + sampling_contexts: tuple[SamplingContext, ...] | None, + stop_rows: tuple[tuple[int, ...], ...] | None, + ) -> PhysicalCohortResult: + return backend.execute_cohort( + input_ids, + max_new_tokens=budgets, + stop_token_ids=shared_stops, + attention_mask=attention_mask, + sampling_contexts=sampling_contexts, + stop_token_ids_by_row=stop_rows, + ) + + @staticmethod + def _physical_results( + selections: Sequence[_PhysicalSelection], + backend_result: PhysicalCohortResult, + ) -> tuple[DriverResult, ...]: + generated_value = getattr(backend_result, "generated_token_ids", None) + if generated_value is None: + raise TypeError( + "physical backend result must expose generated_token_ids" + ) + generated = tuple( + tuple(int(token) for token in row) for row in generated_value + ) + if len(generated) != len(selections): + raise ValueError( + "physical backend returned a different number of generated rows" + ) + traces_value = tuple(getattr(backend_result, "session_traces", ())) + if traces_value and len(traces_value) != len(selections): + raise ValueError( + "physical backend returned a different number of session traces" + ) + + results: list[DriverResult] = [] + for row, selection in enumerate(selections): + tokens = generated[row] + if len(tokens) > selection.request.max_new_tokens: + raise ValueError( + f"physical backend exceeded max_new_tokens for request " + f"{selection.request.request_id!r}" + ) + if traces_value: + trace = traces_value[row] + if not isinstance(trace, SessionTrace): + raise TypeError( + "physical backend session_traces must contain SessionTrace" + ) + else: + stopped = ( + bool(tokens) + and tokens[-1] in selection.request.stop_token_ids + ) + trace = SessionTrace( + request_id=selection.request.request_id, + backend=selection.backend.name, + cache_kind=selection.backend.capabilities.cache_kind, + sampled=selection.request.is_sampled, + emitted=len(tokens), + finish_reason="stop" if stopped else "length", + route_ahead_status=( + "enabled" + if selection.backend.capabilities.supports_route_ahead + else "disabled" + ), + pairing_evidence=selection.backend.capabilities.pairing_evidence, + executor_evidence=selection.backend.capabilities.executor_evidence, + ) + trace.request_id = selection.request.request_id + finish_reason = trace.finish_reason or ( + "stop" + if tokens and tokens[-1] in selection.request.stop_token_ids + else "length" + ) + results.append( + DriverResult( + request_id=selection.request.request_id, + output_token_ids=tokens, + finish_reason=finish_reason, + backend=selection.backend.name, + trace=trace, + rounds=(), + fallback_reason=selection.fallback_reason, + ) + ) + return tuple(results) + + @staticmethod + def _cohort( + selections: Sequence[_Selection], + ) -> tuple[tuple[_Selection, ...], ...]: + grouped: OrderedDict[tuple[int, Hashable], list[_Selection]] = ( + OrderedDict() + ) + for selection in selections: + key = (id(selection.backend), selection.compatibility_key) + grouped.setdefault(key, []).append(selection) + return tuple(tuple(rows) for rows in grouped.values()) + + @staticmethod + def _run_cohort(rows: Sequence[_ActiveSession]) -> None: + while True: + unfinished = [ + row + for row in rows + if not row.selection.backend.is_finished(row.session) + ] + if not unfinished: + return + for row in unfinished: + backend = row.selection.backend + backend.draft(row.session) + round_result = backend.verify(row.session) + output_progress = len(backend.output_token_ids(row.session)) + progress = max(output_progress, round_result.emitted_length) + if ( + not round_result.finished + and progress <= row.observable_progress + ): + raise BackendProgressError( + f"backend {backend.name} made no progress for request " + f"{row.selection.request.request_id!r}" + ) + row.observable_progress = max(row.observable_progress, progress) + row.rounds.append(round_result) + + @staticmethod + def _cleanup( + rows: Sequence[_ActiveSession], *, restore: bool + ) -> tuple[BaseException, ...]: + errors: list[BaseException] = [] + if restore: + for row in rows: + try: + row.selection.backend.restore( + row.session, row.initial_snapshot + ) + except BaseException as exc: + errors.append(exc) + for row in rows: + try: + row.selection.backend.release(row.session) + except BaseException as exc: + errors.append(exc) + return tuple(errors) + + @staticmethod + def _result(item: _ActiveSession) -> DriverResult: + backend = item.selection.backend + trace = backend.trace(item.session) + trace.request_id = item.selection.request.request_id + finish_reason = trace.finish_reason or "length" + return DriverResult( + request_id=item.selection.request.request_id, + output_token_ids=backend.output_token_ids(item.session), + finish_reason=finish_reason, + backend=backend.name, + trace=trace, + rounds=tuple(item.rounds), + fallback_reason=item.selection.fallback_reason, + ) + + +__all__ = [ + "BackendProgressError", + "CohortPlan", + "DriverResult", + "PhysicalCohortDriverResult", + "SessionDriver", + "SessionCleanupError", + "UnsupportedRequestError", +] diff --git a/moe_infinity/utils/config.py b/moe_infinity/utils/config.py index 00652cf3..60abba79 100644 --- a/moe_infinity/utils/config.py +++ b/moe_infinity/utils/config.py @@ -74,6 +74,24 @@ class ArcherConfig: "help": "Enable attention backend offloading. Default False (uses HuggingFace attention)." }, ) + enable_deepseek_mla_paging: bool = field( + default=False, + metadata={ + "help": "Enable experimental batch-one DeepSeek V2/V3 MLA paging. Default False." + }, + ) + max_resident_paged_speculative_sessions: int = field( + default=1, + metadata={ + "help": "Maximum concurrent resident paged-MLA speculative sessions. Default 1; set 0 to force Stage 4a fallback." + }, + ) + min_free_mla_blocks_after_admission: int = field( + default=1, + metadata={ + "help": "Minimum MLA blocks that remain free after reserving all active and newly admitted requests' full declared budgets plus maximum transient DFlash verify peaks. Default 1." + }, + ) enable_kv_cache_offload: bool = field( default=False, metadata={ @@ -160,3 +178,17 @@ def __post_init__(self): raise ValueError( f"device_memory_ratio ({self.device_memory_ratio}) + kv_cache_memory_ratio ({self.kv_cache_memory_ratio}) > 1.0" ) + if ( + type(self.max_resident_paged_speculative_sessions) is not int + or self.max_resident_paged_speculative_sessions < 0 + ): + raise ValueError( + "max_resident_paged_speculative_sessions must be an integer >= 0" + ) + if ( + type(self.min_free_mla_blocks_after_admission) is not int + or self.min_free_mla_blocks_after_admission < 1 + ): + raise ValueError( + "min_free_mla_blocks_after_admission must be an integer >= 1" + ) diff --git a/pyproject.toml b/pyproject.toml index c31eb89c..b7367b48 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,21 @@ [build-system] -requires = ["setuptools>=75.3.2", "wheel", "torch"] +requires = ["setuptools>=75.3.2", "setuptools-scm>=8", "wheel", "torch"] build-backend = "setuptools.build_meta" +[tool.setuptools_scm] +# Version is derived from git tags (e.g. tag "v0.1.0" -> version "0.1.0"). +# The resolved version is written to this file, which moe_infinity/__init__.py +# imports at runtime. The file is git-ignored and regenerated on every build. +version_file = "moe_infinity/_version.py" +# PyPI/PEP 440 forbid local version segments (the "+gHASH" / "+dYYYYMMDD" +# suffix), so strip them. Untagged nightly builds then become PyPI-valid +# pre-releases, e.g. "0.0.2.dev4" for the 4th commit after tag "v0.0.1". +local_scheme = "no-local-version" +# Used only when building with neither git metadata nor sdist PKG-INFO +# (e.g. a GitHub "Download ZIP" tarball). Normal sdist installs read the +# version from the PKG-INFO baked into the tarball, so this is a last resort. +fallback_version = "0.0.0" + [tool.ruff] line-length = 80 diff --git a/setup.py b/setup.py index ea0e4818..18b33ce2 100644 --- a/setup.py +++ b/setup.py @@ -398,7 +398,9 @@ def _find_nvtx_include_dir() -> Optional[str]: # install all files in the package, rather than just the egg setup( name="moe_infinity", - version=os.getenv("MOEINF_VERSION", "0.0.1"), + # version is supplied dynamically by setuptools-scm (see pyproject.toml + # [tool.setuptools_scm]); it is derived from git tags at build time and + # baked into the sdist's PKG-INFO so it survives a source reinstall. packages=find_packages(exclude=["extensions", "extensions.*"]), include_package_data=True, install_requires=install_requires, diff --git a/tests/examples/test_shift.py b/tests/examples/test_shift.py index 3b7e944b..4656cf1d 100644 --- a/tests/examples/test_shift.py +++ b/tests/examples/test_shift.py @@ -164,10 +164,7 @@ def load_single_dataset(dataset, name, split): custom_kwargs = {"decoder_start_token_id": 0} elif "nllb" in args.model_name_or_path.lower(): custom_kwargs = {"forced_bos_token_id": 256057} # translate to French -elif ( - "mixtral" in args.model_name_or_path.lower() - or "snowflake" in args.model_name_or_path.lower() -): +elif "mixtral" in args.model_name_or_path.lower(): custom_kwargs = {"pad_token_id": tokenizer.eos_token_id} elif "deepseek" in args.model_name_or_path.lower(): custom_kwargs = {} diff --git a/tests/python/dflash/test_bare_hf_backend.py b/tests/python/dflash/test_bare_hf_backend.py new file mode 100644 index 00000000..719bf044 --- /dev/null +++ b/tests/python/dflash/test_bare_hf_backend.py @@ -0,0 +1,243 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import pytest +import torch + +from moe_infinity.spec_decode.backends import PhysicalCohortBackend +from moe_infinity.spec_decode.backends_bare_hf import ( + BareHFCohortResult, + BatchedBareHFBackend, +) +from moe_infinity.spec_decode.protocols import ( + NativeStepTrace, + RequestSpec, + SamplingContext, +) +from tests.python.dflash import test_batched_spec as batched + + +def _request( + request_id: str, + prompt: Sequence[int] = batched.PROMPT_A, + *, + budget: int = 8, + temperature: float = 0.0, +) -> RequestSpec: + return RequestSpec( + request_id=request_id, + prompt_token_ids=tuple(prompt), + max_new_tokens=budget, + sampling=SamplingContext(temperature=temperature), + ) + + +def _new_tokens(result: BareHFCohortResult) -> list[list[int]]: + return [list(row) for row in result.generated_token_ids] + + +def test_backend_declares_physical_mixed_bare_hf_capabilities() -> None: + spec, _ = batched._tiny_spec() + backend = BatchedBareHFBackend(spec) + + assert isinstance(backend, PhysicalCohortBackend) + assert backend.name == "dflash-batched-bare-hf" + assert backend.capabilities.supports_batch + assert backend.capabilities.supports_sampling + assert backend.capabilities.supports_ragged_rows + assert backend.capabilities.cache_kind == "dense_dynamic" + assert not backend.capabilities.supports_route_ahead + assert not backend.capabilities.executor_evidence.wiring_reachable + assert not backend.capabilities.supports_rich_forward + assert backend.supports(_request("greedy")) + assert backend.supports(_request("sampled", temperature=0.7)) + assert backend.cohort_key(_request("greedy")) == backend.cohort_key( + _request("sampled", temperature=0.7) + ) + + +def test_backend_direct_entry_preserves_padding_positions_budgets_and_diagnostics() -> ( + None +): + spec, target = batched._tiny_spec() + backend = BatchedBareHFBackend(spec) + prompts = [batched.PROMPT_A, batched.PROMPT_B] + input_ids, attention_mask, _ = batched._left_pad(prompts) + budgets = (6, 14) + + result = backend.execute_cohort( + input_ids, + max_new_tokens=budgets, + stop_token_ids=(), + attention_mask=attention_mask, + ) + + assert result.generated_lengths == budgets + assert _new_tokens(result) == [ + batched._plain_new(target, prompt, budget) + for prompt, budget in zip(prompts, budgets) + ] + assert result.target_cache is not None + assert result.step_trace + assert all( + step.target_cache_len == step.start for step in result.step_trace + ) + + +def test_backend_mixed_accept_lengths_refeeds_without_double_emission( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, target = batched._tiny_spec() + backend = BatchedBareHFBackend(spec) + prompts = [batched.PROMPT_A, batched.PROMPT_B] + pads = [0, len(batched.PROMPT_A) - len(batched.PROMPT_B)] + budget = 21 + streams = batched._greedy_streams(target, prompts, budget) + true_drafts = batched._true_continuation_drafts(streams, pads) + + def draft_fn(start: int, row: int) -> list[int]: + if row == 0: + return true_drafts(start, row) + base = start - pads[row] + wrong = (streams[row][base + 1] + 1) % batched.TINY_VOCAB + return [wrong] * (batched.TINY_BLOCK_SIZE - 1) + + batched._install_scripted_batched_drafter( + monkeypatch, spec, draft_fn, batch=2 + ) + input_ids, attention_mask, _ = batched._left_pad(prompts) + + result = backend.execute_cohort( + input_ids, + max_new_tokens=(budget, budget), + stop_token_ids=(), + attention_mask=attention_mask, + ) + + assert _new_tokens(result) == [ + batched._plain_new(target, prompt, budget) for prompt in prompts + ] + by_step: dict[int, list[NativeStepTrace]] = {} + for step in result.step_trace: + by_step.setdefault(step.prev_start, []).append(step) + shared_steps = [steps for steps in by_step.values() if len(steps) == 2] + assert shared_steps + assert all( + sorted(step.accept for step in steps) + == [0, batched.TINY_BLOCK_SIZE - 1] + for steps in shared_steps + ) + assert all( + steps[0].start - previous == 1 for previous, steps in by_step.items() + ) + + +def test_backend_eos_bonus_finishes_only_its_row( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, target = batched._tiny_spec() + backend = BatchedBareHFBackend(spec) + prompts = [batched.PROMPT_A, batched.PROMPT_B] + pads = [0, len(batched.PROMPT_A) - len(batched.PROMPT_B)] + budget = 20 + streams = batched._greedy_streams(target, prompts, budget) + batched._install_scripted_batched_drafter( + monkeypatch, + spec, + batched._true_continuation_drafts(streams, pads), + batch=2, + ) + batched._force_target_argmax_rows( + monkeypatch, spec, {0: {2: batched.EOS_ID}} + ) + input_ids, attention_mask, _ = batched._left_pad(prompts) + + result = backend.execute_cohort( + input_ids, + max_new_tokens=(budget, budget), + stop_token_ids=(batched.EOS_ID,), + attention_mask=attention_mask, + ) + + assert result.generated_token_ids[0][-1] == batched.EOS_ID + assert result.generated_lengths == (4, budget) + assert list(result.generated_token_ids[1]) == batched._plain_new( + target, prompts[1], budget + ) + + +def test_legacy_adapter_delegates_and_only_adapts_result( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor([batched.PROMPT_A, batched.PROMPT_A]) + target_cache = object() + draft_cache = object() + trace = NativeStepTrace(5, 0, 6, 1, 6, None) + calls: list[dict[str, Any]] = [] + + def execute( + self: BatchedBareHFBackend, + cohort_input_ids: torch.Tensor, + *, + max_new_tokens: tuple[int, ...], + stop_token_ids: tuple[int, ...], + attention_mask: torch.Tensor, + sampling_contexts: tuple[SamplingContext, ...] | None = None, + stop_token_ids_by_row: tuple[tuple[int, ...], ...] | None = None, + ) -> BareHFCohortResult: + calls.append( + { + "backend": self, + "input_ids": cohort_input_ids, + "budgets": max_new_tokens, + "stops": stop_token_ids, + "mask": attention_mask, + "sampling_contexts": sampling_contexts, + "stop_rows": stop_token_ids_by_row, + } + ) + return BareHFCohortResult( + generated_token_ids=((41,), (51, 52)), + step_trace=(trace,), + target_cache=target_cache, + draft_cache=draft_cache, + ) + + monkeypatch.setattr(BatchedBareHFBackend, "execute_cohort", execute) + + output = spec._generate_batched( + input_ids, + max_new_tokens=[1, 2], + stop_token_ids=[7], + attention_mask=None, + ) + + assert len(calls) == 1 + assert calls[0]["budgets"] == (1, 2) + assert calls[0]["stops"] == (7,) + assert torch.equal(calls[0]["mask"], torch.ones_like(input_ids)) + assert calls[0]["sampling_contexts"] is None + assert calls[0]["stop_rows"] is None + assert output[:, input_ids.shape[1] :].tolist() == [[41, 0], [51, 52]] + assert spec.last_generated_lengths == [1, 2] + assert spec.step_trace == [trace] + assert spec.last_target_cache is target_cache + assert spec.last_draft_cache is draft_cache + + +def test_legacy_adapter_rejects_non_binary_attention_mask() -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor([batched.PROMPT_A, batched.PROMPT_A]) + attention_mask = torch.ones_like(input_ids, dtype=torch.float32) + attention_mask[0, 0] = 0.5 + + with pytest.raises(ValueError, match="0/1 valued"): + spec._generate_batched( + input_ids, + max_new_tokens=2, + stop_token_ids=None, + attention_mask=attention_mask, + ) diff --git a/tests/python/dflash/test_batched_spec.py b/tests/python/dflash/test_batched_spec.py index 0000cfe3..1fb1cee7 100644 --- a/tests/python/dflash/test_batched_spec.py +++ b/tests/python/dflash/test_batched_spec.py @@ -20,9 +20,9 @@ Plus the C0 batched pure ops (``acceptance_lengths`` / ``committed_tokens_ragged`` / ``build_block_with_prefixes``) against their -per-row v1 counterparts, and the dispatch guard rails (sampled batch>1 and -MoE-rich batch>1 raise ``NotImplementedError``; right-padded/non-monotone -masks raise ``ValueError``). +per-row v1 counterparts, and the dispatch guard rails (sampled bare-HF +batch>1 uses physical execution while MoE-rich batch>1 uses independent +request sessions; right-padded/non-monotone masks raise ``ValueError``). Trace convention (batched): one ``NativeStepTrace`` per ACTIVE row per step; ``accept`` is that row's effective accept (``cc_b - 1``) and ``start`` the @@ -529,26 +529,43 @@ def test_batch_one_via_batched_path_with_stop_ids_equals_legacy(): # --------------------------------------------------------------------------- -def test_batched_sampled_raises_not_implemented(): +def test_batched_sampled_bare_hf_is_supported(): spec, _ = _tiny_spec() ids = torch.tensor([PROMPT_A, PROMPT_A]) - with pytest.raises(NotImplementedError, match="greedy-only"): - spec.generate(ids, max_new_tokens=4, temperature=0.7) + output = spec.generate( + ids, + max_new_tokens=4, + temperature=0.7, + generator=torch.Generator().manual_seed(17), + ) + assert output.shape == (2, len(PROMPT_A) + 4) + assert spec.last_generated_lengths == [4, 4] -def test_batched_moe_rich_target_raises_not_implemented(): +def test_batched_moe_rich_target_uses_independent_request_sessions(): from moe_infinity.entrypoints.big_modeling import MoE spec, target = _tiny_spec() shell = MoE.__new__(MoE) shell.model = target + shell._configure_hook = lambda _input_ids: None + shell._cached_past_key_values = None + shell._native_attention_backend = None + shell._resolve_native_input_device = lambda: torch.device("cpu") spec = DFlashSpeculator.from_models( shell, spec.draft, config=spec.config, device="cpu" ) ids = torch.tensor([PROMPT_A, PROMPT_A]) - with pytest.raises(NotImplementedError, match="bare HF target"): - spec.generate(ids, max_new_tokens=4) + output = spec.generate(ids, max_new_tokens=[2, 4]) + + assert output.shape == (2, ids.shape[1] + 4) + assert torch.equal(output[:, : ids.shape[1]], ids) + assert spec.last_generated_lengths == [2, 4] + assert [trace.request_id for trace in spec.last_session_traces] == [ + "direct-0", + "direct-1", + ] def test_right_padded_attention_mask_rejected(): diff --git a/tests/python/dflash/test_capability_orthogonality.py b/tests/python/dflash/test_capability_orthogonality.py new file mode 100644 index 00000000..76e8fede --- /dev/null +++ b/tests/python/dflash/test_capability_orthogonality.py @@ -0,0 +1,255 @@ +from __future__ import annotations + +from dataclasses import FrozenInstanceError +from types import SimpleNamespace +from unittest.mock import MagicMock + +import pytest +import torch + +from moe_infinity.distributed.expert_executor import DistributedExpertExecutor +from moe_infinity.spec_decode._route_ahead_ctx import route_ahead_context +from moe_infinity.spec_decode._route_ahead_stats import RouteAheadStats +from moe_infinity.spec_decode.backends import DFlashExecutionBackend +from moe_infinity.spec_decode.dflash import DFlashConfig, validate_pairing +from moe_infinity.spec_decode.protocols import ( + ExecutorEvidence, + PairingEvidence, + SessionTrace, +) +from moe_infinity.utils import ArcherConfig + + +def _valid_pairing( + *, checkpoint_scope: tuple[str, ...] = () +) -> PairingEvidence: + return PairingEvidence( + valid=True, + config_valid=True, + dimensions_valid=True, + vocab_valid=True, + mask_valid=True, + layers_valid=True, + block_valid=True, + module_valid=True, + validated_checkpoint_scope=checkpoint_scope, + ) + + +def _executor(*, prefetcher: object | None = None) -> DistributedExpertExecutor: + config = ArcherConfig.load_from_json( + { + "offload_path": "/tmp/moe-infinity-capability-orthogonality", + "trace_capacity": 16, + "prefetch": False, + "speculative_prefetch": True, + } + ) + executor = DistributedExpertExecutor(config) + executor.set_expert_dispatcher(MagicMock(name="ExpertDispatcher")) + if prefetcher is not None: + executor.set_prefetcher(prefetcher) + return executor + + +def _dispatch(executor: DistributedExpertExecutor) -> None: + mask = torch.tensor([[1, 0, 1], [0, 1, 0]], dtype=torch.bool) + executor.dispatch_local( + 4, + torch.zeros(2, 4), + mask, + mask.to(torch.float32), + router_logits=torch.zeros(2, 3), + ) + + +def test_pairing_and_executor_evidence_are_frozen_and_orthogonal() -> None: + pairing = _valid_pairing( + checkpoint_scope=( + "openai/gpt-oss-120b", + "z-lab/gpt-oss-120b-DFlash", + ) + ) + executor = ExecutorEvidence( + wiring_reachable=False, + fallback_reason="executor_unreachable", + ) + + assert pairing.valid + assert not executor.wiring_reachable + assert not hasattr(pairing, "wiring_reachable") + assert not hasattr(executor, "valid") + with pytest.raises(FrozenInstanceError): + pairing.valid = False # type: ignore[misc] + with pytest.raises(FrozenInstanceError): + executor.wiring_reachable = True # type: ignore[misc] + + +def test_valid_published_pair_with_no_executor_reports_valid_unreachable() -> ( + None +): + pairing = _valid_pairing( + checkpoint_scope=( + "openai/gpt-oss-20b", + "z-lab/gpt-oss-20b-DFlash", + ) + ) + speculator = SimpleNamespace( + moe=SimpleNamespace(_native_model_forward_rich=lambda *_args: None), + pairing_evidence=pairing, + executor_evidence=ExecutorEvidence( + wiring_reachable=False, + fallback_reason="executor_unreachable", + ), + ) + + backend = DFlashExecutionBackend(speculator) + + assert backend.capabilities.pairing_evidence.valid + assert not backend.capabilities.executor_evidence.wiring_reachable + assert not backend.capabilities.supports_route_ahead + + +def test_executor_wiring_does_not_make_an_invalid_pair_valid() -> None: + invalid = DFlashConfig( + block_size=8, + mask_token_id=63, + target_layer_ids=[1, 3], + num_target_layers=4, + hidden_size=16, + vocab_size=64, + ) + wired_target = SimpleNamespace( + hidden_size=32, + vocab_size=64, + num_hidden_layers=4, + expert_executor=object(), + ) + + with pytest.raises(ValueError, match="hidden_size"): + validate_pairing(invalid, wired_target) + + +def test_prefetch_exception_records_fallback_and_preserves_legacy_dispatch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + prefetcher = MagicMock(name="ExpertPrefetcher") + prefetcher.fetch_experts_lock_cache.side_effect = RuntimeError( + "prefetch boom" + ) + executor = _executor(prefetcher=prefetcher) + stats = RouteAheadStats() + stats.begin_step() + + with route_ahead_context(prefetcher=prefetcher, stats=stats): + _dispatch(executor) + + evidence = stats.executor_evidence + assert evidence.wiring_reachable + assert evidence.prefetcher_present + assert evidence.attempted_layers == (4,) + assert evidence.fired_layers == () + assert evidence.actual_expert_union == frozenset({(4, 0), (4, 1), (4, 2)}) + assert evidence.fallback_reason == "prefetch_exception:RuntimeError" + assert executor._pending_prefetch is not None + assert executor._pending_prefetch[3] is not None + assert sorted( + call.args[1] + for call in executor.expert_dispatcher.enqueue_expert.call_args_list + ) == [0, 1, 2] + + +def test_prefetch_exceptions_cannot_change_waited_expert_output( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + prefetcher = MagicMock(name="ExpertPrefetcher") + prefetcher.fetch_experts_lock_cache.side_effect = RuntimeError("lock boom") + prefetcher.correct_prefetch.side_effect = RuntimeError("correct boom") + prefetcher.speculative_prefetch.side_effect = RuntimeError("legacy boom") + executor = _executor(prefetcher=prefetcher) + executor.expert_dispatcher.wait_expert.return_value = "legacy-output" + stats = RouteAheadStats() + stats.begin_step() + + with route_ahead_context(prefetcher=prefetcher, stats=stats): + _dispatch(executor) + result = executor.wait_dispatch_local() + + assert result == "legacy-output" + assert stats.executor_evidence.fallback_reason == ( + "prefetch_exception:RuntimeError" + ) + + +@pytest.mark.parametrize( + ("with_context", "mask", "reason"), + [ + (False, torch.tensor([[1, 0]], dtype=torch.bool), "context_inactive"), + (True, torch.tensor([[1, 0]], dtype=torch.bool), "prefetcher_absent"), + (True, torch.zeros(1, 2, dtype=torch.bool), "empty_actual_union"), + ], +) +def test_capability_misses_fall_back_without_changing_dispatch( + monkeypatch: pytest.MonkeyPatch, + with_context: bool, + mask: torch.Tensor, + reason: str, +) -> None: + monkeypatch.setattr(torch.cuda, "device_count", lambda: 1) + executor = _executor() + stats = RouteAheadStats() + stats.begin_step() + + def run() -> None: + executor.dispatch_local( + 2, + torch.zeros(mask.shape[0], 4), + mask, + mask.to(torch.float32), + router_logits=torch.zeros(mask.shape), + ) + + if with_context: + with route_ahead_context(stats=stats): + run() + evidence = stats.executor_evidence + else: + run() + evidence = executor.last_executor_evidence + + assert evidence.fallback_reason == reason + assert executor._pending_prefetch is not None + assert executor._pending_prefetch[3] is not None + + +def test_common_trace_serializes_pairing_and_executor_evidence_separately() -> ( + None +): + pairing = _valid_pairing() + executor = ExecutorEvidence( + wiring_reachable=True, + prefetcher_present=True, + attempted_layers=(1,), + fired_layers=(1,), + actual_expert_union=frozenset({(1, 3)}), + prefetched_bytes=4096, + coverage=1.0, + wasted_prefetch_bytes=0, + cache_hit_rate=0.75, + ) + trace = SessionTrace( + request_id="req", + backend="native", + cache_kind="dense_dynamic", + sampled=False, + pairing_evidence=pairing, + executor_evidence=executor, + ) + + payload = trace.as_dict() + assert payload["pairing_evidence"]["valid"] is True + assert payload["executor_evidence"]["wiring_reachable"] is True + assert "wiring_reachable" not in payload["pairing_evidence"] + assert "valid" not in payload["executor_evidence"] diff --git a/tests/python/dflash/test_compatibility_matrix.py b/tests/python/dflash/test_compatibility_matrix.py new file mode 100644 index 00000000..dc32c570 --- /dev/null +++ b/tests/python/dflash/test_compatibility_matrix.py @@ -0,0 +1,269 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[3] +DFLASH_DOC = ROOT / "docs" / "dflash.md" +MODEL_MATRIX = ROOT / "docs" / "model-compatibility.md" +SERVING_DOC = ROOT / "docs" / "serving.md" +DESIGN = ( + ROOT + / "docs" + / "superpowers" + / "specs" + / "2026-08-17-dflash-unified-execution-design.md" +) +PLAN = ( + ROOT + / "docs" + / "superpowers" + / "plans" + / "2026-08-17-dflash-unified-execution.md" +) + + +def _text(path: Path) -> str: + return path.read_text(encoding="utf-8") + + +def test_matrix_reports_pairing_and_executor_evidence_separately() -> None: + matrix = _text(MODEL_MATRIX) + assert "DFlash pairing evidence" in matrix + assert "Executor / route-ahead evidence" in matrix + assert "GPT-OSS" in matrix + assert "valid published pairs" in matrix + assert "no executor route-ahead" in matrix + assert "No real DeepSeek DFlash pair" in matrix + + +def test_matrix_gates_rich_batch_and_paged_claims_on_capabilities() -> None: + matrix = _text(MODEL_MATRIX) + assert "Rich execution capability" in matrix + assert "Serving cache capability" in matrix + assert "row-aware capability guard" in matrix + assert "default-off" in matrix + assert "DeepSeek V2/V3" in matrix + assert "batch-1 greedy" in matrix + assert "resident-only; no swap/preemption" in matrix + assert "Qwen/hybrid fallback" in matrix + + +def test_dflash_documentation_records_unified_semantics_and_limits() -> None: + doc = _text(DFLASH_DOC) + required = ( + "one semantic core", + "SessionDriver", + "mixed greedy and sampled", + "per-row generator", + "correlated", + "not bit-exact across batch shapes", + "dense cache reconstruction", + "last_generated_lengths", + "right-padded", + "grouped per-request", + "physically batched", + "temporary_dynamic", + "paged_mla", + "draft cache remains separate", + "cancellation", + "preemption", + ) + for claim in required: + assert claim in doc + + +def test_deprecated_generate_warning_and_serving_fallback_are_documented() -> ( + None +): + combined = _text(DFLASH_DOC) + _text(SERVING_DOC) + assert "MoE.generate() is deprecated" in combined + assert "DeprecationWarning" in combined + assert "sampled serving" in combined + assert "temporary DynamicCache" in combined + assert "fallback" in combined + assert "not evidence of sampled serving" in combined + + +def test_trace_fields_are_stable_for_direct_and_serving() -> None: + doc = _text(DFLASH_DOC) + for field in ( + "request_id", + "backend", + "cache_kind", + "round_count", + "accepted", + "committed", + "emitted", + "rollback", + "replay", + "pairing_evidence", + "executor_evidence", + ): + assert f"`{field}`" in doc + assert "direct and serving" in doc + + +def test_checked_in_design_and_plan_record_task_8_5_dependency_order() -> None: + design = _text(DESIGN) + plan = _text(PLAN) + for text in (design, plan): + assert "Task 8.5" in text + assert "DeepSeek MLA prerequisite" in text + assert "Task 8.5 -> Task 9" in text + assert "Stage 4b is default-off" in design + assert "actual delivered behavior" in plan + + +def test_tiny_benchmark_reports_measured_and_observed_fields() -> None: + command = [ + sys.executable, + "benchmarks/dflash/unified_execution_benchmark.py", + "--fixture", + "tiny", + "--json", + ] + completed = subprocess.run( + command, cwd=ROOT, check=True, capture_output=True, text=True + ) + repeated = subprocess.run( + command, cwd=ROOT, check=True, capture_output=True, text=True + ) + report = json.loads(completed.stdout) + repeated_report = json.loads(repeated.stdout) + required = { + "fixture", + "prefill_latency_ms", + "verify_latency_ms", + "decode_elapsed_seconds", + "decode_committed_tokens_per_second", + "sample_count", + "round_count", + "accepted_drafts", + "committed_tokens", + "rollback_count", + "replay_count", + "rng_order_invariant", + "sampled_tvd_value", + "sampled_kl_value", + "metric_units", + "cache_pages_peak", + "execution_mode", + "pairing_evidence", + "executor_evidence", + "per_request_rich_calls", + "physical_rich_calls", + } + assert required <= report.keys() + assert report["fixture"] == "tiny" + assert report["measurement_scope"] == "synthetic no-checkpoint CPU fixture" + assert report["prefill_latency_ms"] >= 0 + assert report["verify_latency_ms"] >= 0 + assert report["decode_committed_tokens_per_second"] > 0 + assert report["sampled_tvd_value"] >= 0 + assert report["sampled_kl_value"] >= 0 + assert report["sample_count"] == report["round_count"] + assert ( + report["round_count"] + == report["accepted_drafts"] + report["rollback_count"] + ) + assert ( + report["committed_tokens"] + == report["round_count"] + report["accepted_drafts"] + ) + assert report["decode_committed_tokens_per_second"] == ( + report["committed_tokens"] / report["decode_elapsed_seconds"] + ) + assert report["metric_units"] == { + "prefill_latency_ms": "milliseconds per prefill operation", + "verify_latency_ms": "milliseconds per verify operation", + "decode_elapsed_seconds": "seconds", + "decode_committed_tokens_per_second": "committed tokens per second", + "sampled_tvd_value": "dimensionless", + "sampled_kl_value": "nats", + "cache_pages_peak": "pages", + "cancellation_latency_ms": "milliseconds", + } + deterministic_fields = { + "sample_count", + "round_count", + "accepted_drafts", + "committed_tokens", + "rollback_count", + "replay_count", + "rng_order_invariant", + "sampled_tvd_value", + "sampled_kl_value", + "cache_pages_peak", + "execution_mode", + "pairing_evidence", + "executor_evidence", + } + assert {key: report[key] for key in deterministic_fields} == { + key: repeated_report[key] for key in deterministic_fields + } + + +def test_tiny_validation_fails_closed_and_passes_all_local_gates() -> None: + command = [ + sys.executable, + "benchmarks/dflash/validate_unified_execution.py", + "--fixture", + "tiny", + "--require-cache-invariants", + "--require-order-invariance", + "--json", + ] + completed = subprocess.run( + command, cwd=ROOT, check=True, capture_output=True, text=True + ) + report = json.loads(completed.stdout) + assert report["status"] == "PASS" + assert report["checkpoint_downloads"] is False + assert report["cache_invariants"] is True + assert report["ownership_isolation"] is True + assert report["order_invariance"] is True + assert report["required_gpu_fixture"] is False + assert isinstance(report["sampled_tvd_value"], float) + assert isinstance(report["sampled_kl_value"], float) + assert report["sampled_tvd_pass"] is True + assert report["sampled_kl_pass"] is True + + +def test_require_gpu_is_readiness_only_and_fails_when_fixture_env_is_disabled() -> ( + None +): + command = [ + sys.executable, + "benchmarks/dflash/validate_unified_execution.py", + "--fixture", + "tiny", + "--require-gpu", + "--json", + ] + environment = dict(os.environ, MOE_DFLASH_GPU="0") + completed = subprocess.run( + command, + cwd=ROOT, + check=False, + capture_output=True, + text=True, + env=environment, + ) + report = json.loads(completed.stdout) + assert completed.returncode == 1 + assert report["status"] == "FAIL" + assert report["gpu_readiness_required"] is True + assert report["gpu_readiness_pass"] is False + assert report["gpu_harness_executed"] is False + assert report["gpu_gate_kind"] == "readiness only" + + +def test_docs_distinguish_gpu_readiness_from_harness_execution() -> None: + combined = _text(DFLASH_DOC) + _text(PLAN) + assert "--require-gpu is a readiness gate" in combined + assert "does not execute the GPU harness" in combined + assert "actual GPU pytest command remains separate and required" in combined diff --git a/tests/python/dflash/test_engine_wire.py b/tests/python/dflash/test_engine_wire.py index c645a71c..1ecf7514 100644 --- a/tests/python/dflash/test_engine_wire.py +++ b/tests/python/dflash/test_engine_wire.py @@ -13,7 +13,8 @@ Also pinned: omitting the kwarg detaches a previously attached strategy (the kwarg is per-call, never sticky), non-greedy params with a drafter configured -still use the standard path (T1 gate), and batch>1 with a drafter fails loudly. +still use the standard path (T1 gate), and rich batch>1 uses independent +request sessions without pretending the singleton engine is physically batched. """ from __future__ import annotations @@ -259,12 +260,20 @@ def test_non_greedy_with_drafter_uses_standard_path(): assert out.shape[1] > len(PROMPT) -def test_batch_larger_than_one_with_drafter_raises(): - """v1 guardrail: spec decoding is batch==1 only; fail loudly.""" +def test_batch_larger_than_one_with_drafter_preserves_every_row(): + """The facade drives independent rich sessions and preserves row order.""" shell, target = _tiny_moe_shell() spec = _tiny_speculator(shell, target) input_ids = torch.tensor([PROMPT, PROMPT], dtype=torch.long) with warnings.catch_warnings(): warnings.simplefilter("ignore", DeprecationWarning) - with pytest.raises(NotImplementedError, match="batch"): - shell.generate(input_ids, do_sample=False, speculative_draft=spec) + output = shell.generate( + input_ids, + do_sample=False, + max_new_tokens=[2, 4], + speculative_draft=spec, + ) + + assert output.shape == (2, len(PROMPT) + 4) + assert torch.equal(output[:, : len(PROMPT)], input_ids) + assert spec.last_generated_lengths == [2, 4] diff --git a/tests/python/dflash/test_loop_retirement.py b/tests/python/dflash/test_loop_retirement.py new file mode 100644 index 00000000..a38c8c6b --- /dev/null +++ b/tests/python/dflash/test_loop_retirement.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +import ast +import inspect +import textwrap +from dataclasses import dataclass, field +from types import SimpleNamespace +from typing import Any, Hashable + +import pytest +import torch + +from moe_infinity.spec_decode.backends_bare_hf import BareHFCohortResult +from moe_infinity.spec_decode.dflash import DFlashSpeculator +from moe_infinity.spec_decode.protocols import ( + BackendCapabilities, + NativeStepTrace, + RequestSpec, + SamplingContext, + SessionTrace, +) +from moe_infinity.spec_decode.session_driver import ( + DriverResult, + SessionDriver, + UnsupportedRequestError, +) +from tests.python.dflash import test_batched_spec as batched + + +def _call_tail(call: ast.Call) -> str: + function = call.func + if isinstance(function, ast.Name): + return function.id + if isinstance(function, ast.Attribute): + return function.attr + return "" + + +def test_legacy_batch_adapter_has_no_semantic_decode_path() -> None: + tree = ast.parse( + textwrap.dedent(inspect.getsource(DFlashSpeculator._generate_batched)) + ) + calls = { + _call_tail(node) + for node in ast.walk(tree) + if isinstance(node, ast.Call) + } + forbidden = { + "acceptance_lengths", + "acceptance_sampled", + "committed_tokens_ragged", + "committed_tokens_sampled", + "warped_probs", + "_forward_target", + "_verify_target_block", + "_run_drafter", + "_snapshot_target_cache", + "_rollback_target_cache", + "commit_step", + "execute_cohort", + } + + assert not any(isinstance(node, ast.While) for node in ast.walk(tree)) + assert calls.isdisjoint(forbidden) + assert {"RequestSpec", "BatchedBareHFBackend", "SessionDriver"} <= calls + assert "run_physical_cohort" in calls + + +def _capabilities(*, sampling: bool = True) -> BackendCapabilities: + return BackendCapabilities( + supports_batch=True, + supports_sampling=sampling, + supports_ragged_rows=True, + cache_kind="dense_dynamic", + supports_route_ahead=False, + supports_rich_forward=False, + ) + + +def _request( + request_id: str, + *, + prompt: tuple[int, ...] = (1, 2, 3), + budget: int = 2, + sampling: SamplingContext | None = None, + stops: frozenset[int] = frozenset(), +) -> RequestSpec: + return RequestSpec( + request_id=request_id, + prompt_token_ids=prompt, + max_new_tokens=budget, + sampling=sampling or SamplingContext(), + stop_token_ids=stops, + ) + + +@dataclass +class _PhysicalBackend: + name: str + events: list[tuple[str, str]] + accepts: frozenset[str] | None = None + sampling: bool = True + malformed: bool = False + failure: BaseException | None = None + capabilities: BackendCapabilities = field(init=False) + execute_calls: list[dict[str, Any]] = field( + default_factory=list, init=False + ) + + def __post_init__(self) -> None: + self.capabilities = _capabilities(sampling=self.sampling) + + def supports(self, request: RequestSpec) -> bool: + self.events.append(("supports", request.request_id)) + return self.accepts is None or request.request_id in self.accepts + + def cohort_key(self, request: RequestSpec) -> Hashable: + del request + return "dense-mixed" + + def execute_cohort(self, input_ids: torch.Tensor, **kwargs: Any) -> Any: + self.events.append(("execute", self.name)) + self.execute_calls.append(dict(kwargs)) + if self.failure is not None: + raise self.failure + rows = int(input_ids.shape[0]) - int(self.malformed) + generated = tuple((30 + row,) for row in range(rows)) + sampling_contexts = kwargs.get( + "sampling_contexts", + tuple(SamplingContext() for _ in range(int(input_ids.shape[0]))), + ) + if sampling_contexts is None: + sampling_contexts = tuple( + SamplingContext() for _ in range(int(input_ids.shape[0])) + ) + traces = tuple( + SessionTrace( + request_id=f"backend-{row}", + backend=self.name, + cache_kind="dense_dynamic", + sampled=sampling_contexts[row].is_sampled, + emitted=1, + finish_reason="length", + ) + for row in range(rows) + ) + return SimpleNamespace( + generated_token_ids=generated, + generated_lengths=tuple(len(row) for row in generated), + session_traces=traces, + step_trace=(), + target_cache=None, + draft_cache=None, + ) + + +def test_physical_driver_selects_every_row_before_one_execution() -> None: + events: list[tuple[str, str]] = [] + backend = _PhysicalBackend("physical", events) + sampled = SamplingContext( + temperature=0.7, generator=torch.Generator().manual_seed(9) + ) + requests = ( + _request("r0", budget=1, stops=frozenset({8})), + _request("r1", budget=1, sampling=sampled, stops=frozenset({9})), + ) + driver = SessionDriver([backend]) # type: ignore[list-item] + + run = driver.run_physical_cohort( + torch.tensor([[1, 2, 3], [0, 2, 3]]), + requests=requests, + attention_mask=torch.tensor([[1, 1, 1], [0, 1, 1]]), + ) + + assert events == [ + ("supports", "r0"), + ("supports", "r1"), + ("execute", "physical"), + ] + assert run.backend_result.generated_token_ids == ((30,), (31,)) + assert [result.request_id for result in run.results] == ["r0", "r1"] + assert [result.output_token_ids for result in run.results] == [(30,), (31,)] + assert [result.trace.request_id for result in run.results] == ["r0", "r1"] + assert [result.trace.sampled for result in run.results] == [False, True] + assert driver.last_results == run.results + assert driver.last_cohorts[0].row_indices == (0, 1) + + +def test_physical_driver_never_downgrades_sampling() -> None: + events: list[tuple[str, str]] = [] + greedy = _PhysicalBackend("greedy", events, sampling=False) + sampled = _PhysicalBackend("sampled", events) + request = _request("sampled-row", sampling=SamplingContext(temperature=0.8)) + + run = SessionDriver([greedy, sampled]).run_physical_cohort( # type: ignore[list-item] + torch.tensor([[1, 2, 3]]), + requests=(request,), + attention_mask=torch.ones(1, 3, dtype=torch.long), + ) + + assert run.results[0].backend == "sampled" + assert ("execute", "greedy") not in events + + +def test_physical_driver_passes_optional_metadata_in_one_backend_call() -> None: + events: list[tuple[str, str]] = [] + backend = _PhysicalBackend("physical", events) + + SessionDriver([backend]).run_physical_cohort( # type: ignore[list-item] + torch.tensor([[1, 2, 3]]), + requests=(_request("r0", budget=1),), + attention_mask=torch.ones(1, 3, dtype=torch.long), + ) + + assert len(backend.execute_calls) == 1 + assert backend.execute_calls[0]["sampling_contexts"] is None + assert backend.execute_calls[0]["stop_token_ids_by_row"] is None + + +def test_physical_driver_rejects_split_cohort_before_output() -> None: + events: list[tuple[str, str]] = [] + first = _PhysicalBackend("first", events, accepts=frozenset({"r0"})) + second = _PhysicalBackend("second", events, accepts=frozenset({"r1"})) + driver = SessionDriver([first, second]) # type: ignore[list-item] + + with pytest.raises(UnsupportedRequestError, match="one physical cohort"): + driver.run_physical_cohort( + torch.tensor([[1, 2, 3], [1, 2, 3]]), + requests=(_request("r0"), _request("r1")), + attention_mask=torch.ones(2, 3, dtype=torch.long), + ) + + assert not any(event[0] == "execute" for event in events) + assert driver.last_results == () + + +@pytest.mark.parametrize("mode", ["malformed", "failure"]) +def test_physical_driver_publishes_nothing_on_failure(mode: str) -> None: + events: list[tuple[str, str]] = [] + backend = _PhysicalBackend( + "physical", + events, + malformed=mode == "malformed", + failure=RuntimeError("verify failed") if mode == "failure" else None, + ) + driver = SessionDriver([backend]) # type: ignore[list-item] + + with pytest.raises((RuntimeError, ValueError)): + driver.run_physical_cohort( + torch.tensor([[1, 2, 3], [1, 2, 3]]), + requests=(_request("r0"), _request("r1")), + attention_mask=torch.ones(2, 3, dtype=torch.long), + ) + + assert driver.last_results == () + + +def test_physical_driver_clears_published_state_before_normalization() -> None: + events: list[tuple[str, str]] = [] + backend = _PhysicalBackend("physical", events) + driver = SessionDriver([backend]) # type: ignore[list-item] + input_ids = torch.tensor([[1, 2, 3]]) + attention_mask = torch.ones_like(input_ids) + driver.run_physical_cohort( + input_ids, + requests=(_request("valid", budget=1),), + attention_mask=attention_mask, + ) + assert driver.last_results + assert driver.last_cohorts + + with pytest.raises(TypeError, match="RequestSpec"): + driver.run_physical_cohort( + input_ids, + requests=(object(),), # pyright: ignore[reportArgumentType] + attention_mask=attention_mask, + ) + + assert driver.last_results == () + assert driver.last_cohorts == () + + +def test_legacy_adapter_preserves_requests_results_traces_and_diagnostics( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor( + [[0, 0, *batched.PROMPT_B], batched.PROMPT_A], dtype=torch.int32 + ) + attention_mask = torch.tensor([[0, 0, 1, 1, 1], [1, 1, 1, 1, 1]]) + sampling = ( + SamplingContext(), + SamplingContext( + temperature=0.8, + top_k=4, + top_p=0.9, + generator=torch.Generator().manual_seed(17), + ), + ) + target_cache = object() + draft_cache = object() + step = NativeStepTrace(5, 0, 6, 1, 6, None) + traces = ( + SessionTrace( + "direct-0", + "physical", + "dense_dynamic", + False, + emitted=1, + finish_reason="stop", + ), + SessionTrace( + "direct-1", + "physical", + "dense_dynamic", + True, + emitted=2, + finish_reason="length", + ), + ) + backend_result = BareHFCohortResult( + generated_token_ids=((41,), (51, 52)), + step_trace=(step,), + target_cache=target_cache, + draft_cache=draft_cache, + session_traces=traces, + ) + captured: dict[str, Any] = {} + + def run_physical( + self: SessionDriver, + cohort_input_ids: torch.Tensor, + *, + requests: tuple[RequestSpec, ...], + attention_mask: torch.Tensor, + ) -> Any: + captured.update( + driver=self, + input_ids=cohort_input_ids, + requests=requests, + attention_mask=attention_mask, + ) + results = tuple( + DriverResult( + request_id=request.request_id, + output_token_ids=backend_result.generated_token_ids[row], + finish_reason=traces[row].finish_reason or "length", + backend="physical", + trace=traces[row], + rounds=(), + ) + for row, request in enumerate(requests) + ) + return SimpleNamespace(results=results, backend_result=backend_result) + + monkeypatch.setattr( + SessionDriver, "run_physical_cohort", run_physical, raising=False + ) + + output = spec._generate_batched( + input_ids, + max_new_tokens=[1, 2], + stop_token_ids=None, + attention_mask=attention_mask, + sampling_contexts=sampling, + stop_token_ids_by_row=((41,), ()), + ) + + requests = captured["requests"] + assert [request.prompt_token_ids for request in requests] == [ + tuple(batched.PROMPT_B), + tuple(batched.PROMPT_A), + ] + assert [request.max_new_tokens for request in requests] == [1, 2] + assert [request.stop_token_ids for request in requests] == [ + frozenset({41}), + frozenset(), + ] + assert [request.sampling for request in requests] == list(sampling) + assert torch.equal(captured["attention_mask"], attention_mask) + assert output[:, input_ids.shape[1] :].tolist() == [[41, 0], [51, 52]] + assert output.dtype == input_ids.dtype + assert spec.last_generated_lengths == [1, 2] + assert spec.step_trace == [step] + assert spec.last_target_cache is target_cache + assert spec.last_draft_cache is draft_cache + assert spec.last_session_results == captured[ + "driver" + ].last_results or tuple( + result.request_id for result in spec.last_session_results + ) == ("direct-0", "direct-1") + assert spec.last_session_traces == traces diff --git a/tests/python/dflash/test_mixed_sampling_batch.py b/tests/python/dflash/test_mixed_sampling_batch.py new file mode 100644 index 00000000..a41c2b1f --- /dev/null +++ b/tests/python/dflash/test_mixed_sampling_batch.py @@ -0,0 +1,454 @@ +from __future__ import annotations + +from collections.abc import Sequence +from typing import Any + +import pytest +import torch + +import moe_infinity.spec_decode.backends_bare_hf as bare_backend_module +from moe_infinity.spec_decode.backends_bare_hf import BatchedBareHFBackend +from moe_infinity.spec_decode.protocols import RequestSpec, SamplingContext +from tests.python.dflash import test_batched_spec as batched +from tests.python.dflash import test_sampled_spec as sampled + + +def _request( + request_id: str, + *, + temperature: float, + top_k: int = 0, + top_p: float = 1.0, +) -> RequestSpec: + return RequestSpec( + request_id=request_id, + prompt_token_ids=tuple(batched.PROMPT_A), + max_new_tokens=8, + sampling=SamplingContext( + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=torch.Generator().manual_seed(17), + ), + ) + + +def _new_tokens( + output: torch.Tensor, spec: Any, prompt_width: int +) -> list[list[int]]: + lengths = spec.last_generated_lengths + assert lengths is not None + return [ + output[row, prompt_width : prompt_width + length].tolist() + for row, length in enumerate(lengths) + ] + + +def test_bare_hf_backend_advertises_sampled_and_mixed_physical_batches() -> ( + None +): + spec, _ = batched._tiny_spec() + backend = BatchedBareHFBackend(spec) + + greedy = _request("greedy", temperature=0.0) + sampled = _request("sampled", temperature=0.8, top_k=7, top_p=0.9) + + assert backend.capabilities.supports_batch + assert backend.capabilities.supports_sampling + assert backend.supports(greedy) + assert backend.supports(sampled) + assert backend.cohort_key(greedy) == backend.cohort_key(sampled) + + +@pytest.mark.parametrize( + ("argument", "value"), + [ + ("temperature", [0.0]), + ("top_k", [0]), + ("top_p", [1.0]), + ("generator", [torch.Generator().manual_seed(1)]), + ], +) +def test_per_row_sampling_lengths_are_validated_before_prefill( + monkeypatch: pytest.MonkeyPatch, + argument: str, + value: Sequence[Any], +) -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor([batched.PROMPT_A, batched.PROMPT_A]) + called = False + + def forbidden_prefill(*args: Any, **kwargs: Any) -> Any: + nonlocal called + called = True + raise AssertionError("prefill must not run") + + monkeypatch.setattr(spec, "_forward_target", forbidden_prefill) + kwargs: dict[str, Any] = { + "temperature": [0.0, 0.8], + "top_k": [0, 5], + "top_p": [1.0, 0.9], + "generator": [None, torch.Generator().manual_seed(2)], + } + kwargs[argument] = value + + with pytest.raises(ValueError, match=f"{argument}.*batch size 2"): + spec.generate(input_ids, max_new_tokens=4, **kwargs) + + assert not called + + +def test_same_explicit_generator_for_sampled_rows_rejected_before_prefill( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor([batched.PROMPT_A, batched.PROMPT_A]) + generator = torch.Generator().manual_seed(31) + initial_state = generator.get_state().clone() + called = False + + def forbidden_prefill(*args: Any, **kwargs: Any) -> Any: + nonlocal called + called = True + raise AssertionError("prefill must not run") + + monkeypatch.setattr(spec, "_forward_target", forbidden_prefill) + with pytest.raises(ValueError, match="same explicit generator object"): + spec.generate( + input_ids, + max_new_tokens=4, + temperature=[0.7, 0.9], + generator=[generator, generator], + ) + + assert not called + assert torch.equal(generator.get_state(), initial_state) + + +def test_mixed_rows_use_one_physical_target_verify_and_keep_row_policies( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, target = batched._tiny_spec() + prompts = [batched.PROMPT_A, batched.PROMPT_B] + input_ids, attention_mask, width = batched._left_pad(prompts) + verify_batches: list[int] = [] + original_verify = spec._verify_target_block + + def wrapped_verify( + block: torch.Tensor, target_kv: Any, **kwargs: Any + ) -> Any: + verify_batches.append(int(block.shape[0])) + return original_verify(block, target_kv, **kwargs) + + monkeypatch.setattr(spec, "_verify_target_block", wrapped_verify) + sampled_generator = torch.Generator().manual_seed(91) + output = spec.generate( + input_ids, + max_new_tokens=[8, 11], + temperature=[0.0, 0.75], + top_k=[0, 5], + top_p=[1.0, 0.85], + generator=[None, sampled_generator], + attention_mask=attention_mask, + ) + rows = _new_tokens(output, spec, width) + + assert verify_batches and all( + batch_size == 2 for batch_size in verify_batches + ) + assert rows[0] == batched._plain_new(target, prompts[0], 8) + assert len(rows[1]) == 11 + assert rows[0] != rows[1][: len(rows[0])] + + +def test_scalar_sampling_generator_is_cloned_per_sampled_row() -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor([batched.PROMPT_A, batched.PROMPT_A]) + generator = torch.Generator().manual_seed(123) + initial_state = generator.get_state().clone() + + first = spec.generate( + input_ids, + max_new_tokens=12, + temperature=0.8, + top_p=0.9, + generator=generator, + ) + second = spec.generate( + input_ids, + max_new_tokens=12, + temperature=0.8, + top_p=0.9, + generator=generator, + ) + + assert torch.equal(first, second) + assert torch.equal(generator.get_state(), initial_state) + assert first[0].tolist() == first[1].tolist() + + +def test_omitted_batch_generators_become_independent_request_streams( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor([batched.PROMPT_A, batched.PROMPT_A]) + captured: list[SamplingContext] = [] + original = BatchedBareHFBackend.execute_cohort + + def wrapped( + self: BatchedBareHFBackend, + cohort_input_ids: torch.Tensor, + **kwargs: Any, + ) -> Any: + captured.extend(kwargs["sampling_contexts"]) + return original(self, cohort_input_ids, **kwargs) + + monkeypatch.setattr(BatchedBareHFBackend, "execute_cohort", wrapped) + spec.generate(input_ids, max_new_tokens=4, temperature=0.8) + + generators = [context.generator for context in captured] + assert all( + isinstance(generator, torch.Generator) for generator in generators + ) + assert len({id(generator) for generator in generators}) == len(generators) + + +def _run_named_rows(names: Sequence[str]) -> dict[str, list[int]]: + prompts = { + "a": batched.PROMPT_A, + "b": batched.PROMPT_B, + "c": batched.PROMPT_C, + } + policies = { + "a": (0.8, 0, 0.9, 101), + "b": (0.0, 0, 1.0, 202), + "c": (1.1, 7, 1.0, 303), + } + spec, _ = batched._tiny_spec() + input_ids, attention_mask, width = batched._left_pad( + [prompts[name] for name in names] + ) + output = spec.generate( + input_ids, + max_new_tokens=[18] * len(names), + temperature=[policies[name][0] for name in names], + top_k=[policies[name][1] for name in names], + top_p=[policies[name][2] for name in names], + generator=[ + torch.Generator().manual_seed(policies[name][3]) for name in names + ], + attention_mask=attention_mask, + ) + rows = _new_tokens(output, spec, width) + return {name: rows[row] for row, name in enumerate(names)} + + +def test_row_order_and_unrelated_composition_do_not_change_request_stream() -> ( + None +): + forward = _run_named_rows(["a", "b", "c"]) + reverse = _run_named_rows(["c", "b", "a"]) + composed = _run_named_rows(["a", "c"]) + + assert forward == reverse + assert forward["a"] == composed["a"] + assert forward["c"] == composed["c"] + + +def test_sampled_backend_retains_every_slot_proposal_and_row_warp( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, _ = batched._tiny_spec() + prompts = [batched.PROMPT_A, batched.PROMPT_B] + input_ids, attention_mask, _ = batched._left_pad(prompts) + backend = BatchedBareHFBackend(spec) + seen_proposals: list[torch.Tensor] = [] + seen_warps: list[tuple[float, int, float]] = [] + original_acceptance = bare_backend_module.acceptance_sampled + original_warp = bare_backend_module.warped_probs + + def capture_warp( + logits: torch.Tensor, + temperature: float = 1.0, + top_k: int = 0, + top_p: float = 1.0, + ) -> torch.Tensor: + seen_warps.append((temperature, top_k, top_p)) + return original_warp(logits, temperature, top_k, top_p) + + def capture_acceptance( + draft_probs: torch.Tensor, + target_probs: torch.Tensor, + drafts: torch.Tensor, + generator: torch.Generator | None = None, + ) -> Any: + seen_proposals.append(draft_probs.clone()) + return original_acceptance( + draft_probs, target_probs, drafts, generator=generator + ) + + monkeypatch.setattr(bare_backend_module, "warped_probs", capture_warp) + monkeypatch.setattr( + bare_backend_module, "acceptance_sampled", capture_acceptance + ) + backend.execute_cohort( + input_ids, + max_new_tokens=(12, 12), + stop_token_ids=(), + attention_mask=attention_mask, + sampling_contexts=( + SamplingContext( + temperature=0.65, + top_k=4, + top_p=0.8, + generator=torch.Generator().manual_seed(11), + ), + SamplingContext( + temperature=1.2, + top_k=9, + top_p=0.95, + generator=torch.Generator().manual_seed(22), + ), + ), + ) + + assert seen_proposals + assert all( + proposal.shape == (batched.TINY_BLOCK_SIZE - 1, batched.TINY_VOCAB) + for proposal in seen_proposals + ) + assert all( + torch.allclose( + proposal.sum(dim=-1), + torch.ones(batched.TINY_BLOCK_SIZE - 1), + ) + for proposal in seen_proposals + ) + assert (0.65, 4, 0.8) in seen_warps + assert (1.2, 9, 0.95) in seen_warps + + +def test_budget_zero_and_greedy_rows_consume_no_request_rng() -> None: + spec, _ = batched._tiny_spec() + input_ids = torch.tensor([batched.PROMPT_A] * 3) + zero = torch.Generator().manual_seed(1) + greedy = torch.Generator().manual_seed(2) + sampled_generator = torch.Generator().manual_seed(3) + zero_before = zero.get_state().clone() + greedy_before = greedy.get_state().clone() + sampled_before = sampled_generator.get_state().clone() + + spec.generate( + input_ids, + max_new_tokens=[0, 8, 8], + temperature=[0.8, 0.0, 0.8], + generator=[zero, greedy, sampled_generator], + ) + + assert torch.equal(zero_before, zero.get_state()) + assert torch.equal(greedy_before, greedy.get_state()) + assert not torch.equal(sampled_before, sampled_generator.get_state()) + + +def test_per_row_stop_sets_and_budgets_are_independent() -> None: + spec, target = batched._tiny_spec() + prompts = [batched.PROMPT_A, batched.PROMPT_B] + input_ids, attention_mask, width = batched._left_pad(prompts) + first = batched._plain_new(target, prompts[0], 1)[0] + + output = spec.generate( + input_ids, + max_new_tokens=[8, 13], + temperature=[0.0, 0.0], + stop_token_ids=[[first], []], + attention_mask=attention_mask, + ) + rows = _new_tokens(output, spec, width) + + assert rows[0] == [first] + assert rows[1] == batched._plain_new(target, prompts[1], 13) + assert spec.last_generated_lengths == [1, 13] + + +def test_sampled_batch_distribution_matches_looped_singletons_and_plain_target() -> ( + None +): + spec, target = sampled._build_spec( + vocab_size=sampled.PARITY_VOCAB, + mask_token_id=sampled.PARITY_MASK_ID, + ) + input_ids = sampled.PROMPT.repeat(2, 1) + batched_tokens: list[list[int]] = [] + singleton_tokens: list[list[int]] = [] + plain_tokens: list[list[int]] = [] + runs = 120 + for run in range(runs): + output = spec.generate( + input_ids, + max_new_tokens=sampled.PARITY_MAX_NEW, + temperature=0.9, + top_k=5, + top_p=0.9, + generator=[ + torch.Generator().manual_seed(10_000 + 2 * run), + torch.Generator().manual_seed(10_001 + 2 * run), + ], + ) + batched_tokens.extend(_new_tokens(output, spec, sampled.PROMPT_LEN)) + torch.manual_seed(20_000 + run) + singleton = spec.generate( + sampled.PROMPT, + max_new_tokens=sampled.PARITY_MAX_NEW, + temperature=0.9, + top_k=5, + top_p=0.9, + ) + singleton_tokens.append(singleton[0, sampled.PROMPT_LEN :].tolist()) + torch.manual_seed(30_000 + run) + plain_tokens.append( + sampled._plain_sampled_decode( + target, + sampled.PROMPT, + sampled.PARITY_MAX_NEW, + 0.9, + top_k=5, + top_p=0.9, + ) + ) + + batch_hist = sampled._pooled_histogram(batched_tokens, sampled.PARITY_VOCAB) + singleton_hist = sampled._pooled_histogram( + singleton_tokens, sampled.PARITY_VOCAB + ) + plain_hist = sampled._pooled_histogram(plain_tokens, sampled.PARITY_VOCAB) + assert sampled._tvd(batch_hist, singleton_hist) <= 0.13 + assert sampled._tvd(batch_hist, plain_hist) <= 0.13 + assert sampled._kl(batch_hist, singleton_hist) <= 0.08 + assert sampled._kl(batch_hist, plain_hist) <= 0.08 + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable") +def test_cuda_generator_device_mismatch_is_rejected_before_prefill( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, _ = batched._tiny_spec() + spec.device = "cuda" + called = False + + def forbidden_prefill(*args: Any, **kwargs: Any) -> Any: + nonlocal called + called = True + raise AssertionError("prefill must not run") + + monkeypatch.setattr(spec, "_forward_target", forbidden_prefill) + with pytest.raises(ValueError, match="generator device cpu.*cuda"): + spec.generate( + torch.tensor([batched.PROMPT_A, batched.PROMPT_A]), + max_new_tokens=4, + temperature=[0.8, 0.8], + generator=[ + torch.Generator().manual_seed(1), + torch.Generator().manual_seed(2), + ], + ) + assert not called diff --git a/tests/python/dflash/test_public_api_compat.py b/tests/python/dflash/test_public_api_compat.py new file mode 100644 index 00000000..6171fdbd --- /dev/null +++ b/tests/python/dflash/test_public_api_compat.py @@ -0,0 +1,290 @@ +from __future__ import annotations + +import warnings +from types import SimpleNamespace +from typing import Any, cast + +import pytest +import torch + +from moe_infinity.entrypoints.big_modeling import MoE +from moe_infinity.spec_decode import SessionTrace +from moe_infinity.spec_decode.dflash import _normalize_stop_rows +from tests.python.dflash import test_batched_spec as batched +from tests.python.dflash.test_engine_wire import ( + PROMPT, + _tiny_moe_shell, + _tiny_speculator, +) + +WARNING_TEXT = ( + "MoE.generate() is deprecated. Use MoE.serve() for continuous batching " + "with higher throughput. MoE.generate() will be removed in a future version." +) + + +@pytest.mark.parametrize( + ("value", "expected", "per_row"), + [ + (17, ((17,), (17,)), False), + ([17, 19], ((17, 19), (17, 19)), False), + ([[17], [19, 23]], ((17,), (19, 23)), True), + ], +) +def test_stop_rows_accept_scalar_shared_and_per_row_forms( + value: object, + expected: tuple[tuple[int, ...], ...], + per_row: bool, +) -> None: + target = SimpleNamespace(config=SimpleNamespace(eos_token_id=2)) + + rows, actual_per_row = _normalize_stop_rows(target, value, batch=2) + + assert rows == expected + assert actual_per_row is per_row + + +@pytest.mark.parametrize( + "value", + [True, [1, [2]], [[1], 2], [None, 2]], +) +def test_stop_rows_reject_malformed_or_boolean_values(value: object) -> None: + target = SimpleNamespace(config=SimpleNamespace(eos_token_id=2)) + + with pytest.raises(ValueError, match="stop_token_ids"): + _normalize_stop_rows(target, value, batch=2) + + +def test_direct_bare_hf_batch_accepts_scalar_stop_id() -> None: + spec, target = batched._tiny_spec() + stop_id = batched._plain_new(target, batched.PROMPT_A, 1)[0] + input_ids = torch.tensor([batched.PROMPT_A, batched.PROMPT_A]) + + output = spec.generate( + input_ids, + max_new_tokens=8, + stop_token_ids=stop_id, + ) + + assert spec.last_generated_lengths == [1, 1] + assert output[:, input_ids.shape[1]].tolist() == [stop_id, stop_id] + + +def test_moe_generate_preserves_deprecation_warning_text_and_stacklevel() -> ( + None +): + shell, _ = _tiny_moe_shell() + + with warnings.catch_warnings(record=True) as caught: + warnings.simplefilter("always", DeprecationWarning) + shell.generate( + cast(torch.LongTensor, torch.tensor([PROMPT])), + do_sample=False, + max_new_tokens=1, + ) + + warning = next( + item for item in caught if item.category is DeprecationWarning + ) + assert str(warning.message) == WARNING_TEXT + assert warning.filename == __file__ + + +def test_direct_rich_batch_uses_independent_sessions_and_shared_trace_schema( + monkeypatch: pytest.MonkeyPatch, +) -> None: + shell, target = _tiny_moe_shell() + spec = _tiny_speculator(shell, target) + begin_shapes: list[tuple[int, ...]] = [] + real_begin = spec.begin_session + + def begin(input_ids: torch.Tensor, **kwargs: Any): + begin_shapes.append(tuple(input_ids.shape)) + return real_begin(input_ids, **kwargs) + + monkeypatch.setattr(spec, "begin_session", begin) + input_ids = torch.tensor([PROMPT, PROMPT], dtype=torch.long) + + output = spec.generate(input_ids, max_new_tokens=[2, 5]) + + assert begin_shapes == [(1, len(PROMPT)), (1, len(PROMPT))] + assert output.shape == (2, len(PROMPT) + 5) + assert torch.equal(output[:, : len(PROMPT)], input_ids) + assert spec.last_generated_lengths == [2, 5] + assert len(spec.last_session_traces) == 2 + assert all( + isinstance(trace, SessionTrace) for trace in spec.last_session_traces + ) + assert [trace.request_id for trace in spec.last_session_traces] == [ + "direct-0", + "direct-1", + ] + assert all( + trace.backend == "dflash-per-request" + for trace in spec.last_session_traces + ) + + +def test_moe_generate_rich_batch_delegates_once_without_engine_batch( + monkeypatch: pytest.MonkeyPatch, +) -> None: + shell, target = _tiny_moe_shell() + spec = _tiny_speculator(shell, target) + input_ids = torch.tensor([PROMPT, PROMPT], dtype=torch.long) + expected = torch.cat( + [input_ids, torch.tensor([[31, 0, 0], [41, 42, 43]])], dim=1 + ) + calls: list[dict[str, Any]] = [] + + def direct_generate(ids: torch.Tensor, **kwargs: Any) -> torch.Tensor: + calls.append({"input_ids": ids, **kwargs}) + spec.last_generated_lengths = [1, 3] + return expected + + monkeypatch.setattr(spec, "generate", direct_generate) + engine_calls: list[dict[str, Any]] = [] + real_engine_generate = shell._native_generation_engine.generate + + def engine_generate(**kwargs: Any): + engine_calls.append(kwargs) + return real_engine_generate(**kwargs) + + monkeypatch.setattr( + shell._native_generation_engine, "generate", engine_generate + ) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + actual = shell.generate( + cast(torch.LongTensor, input_ids), + do_sample=False, + max_new_tokens=[1, 3], + speculative_draft=spec, + ) + + assert torch.equal(actual, expected) + assert len(calls) == 1 + assert calls[0]["max_new_tokens"] == [1, 3] + assert engine_calls == [] + assert shell._native_generation_engine.spec_strategy is None + + +def test_real_moe_rich_batch_accepts_scalar_eos_and_detaches_strategy() -> None: + shell, target = _tiny_moe_shell() + spec = _tiny_speculator(shell, target) + input_ids = torch.tensor([PROMPT, PROMPT], dtype=torch.long) + probe = spec.generate(input_ids[:1], max_new_tokens=1) + stop_id = int(probe[0, len(PROMPT)]) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + output = shell.generate( + cast(torch.LongTensor, input_ids), + do_sample=False, + max_new_tokens=8, + eos_token_id=stop_id, + speculative_draft=spec, + ) + + assert output.shape == (2, len(PROMPT) + 1) + assert output[:, len(PROMPT)].tolist() == [stop_id, stop_id] + assert spec.last_generated_lengths == [1, 1] + assert shell._native_generation_engine.spec_strategy is None + + +def test_moe_rich_batch_detaches_strategy_when_speculator_fails( + monkeypatch: pytest.MonkeyPatch, +) -> None: + shell, target = _tiny_moe_shell() + spec = _tiny_speculator(shell, target) + input_ids = torch.tensor([PROMPT, PROMPT], dtype=torch.long) + + def fail(*_args: Any, **_kwargs: Any) -> torch.Tensor: + raise RuntimeError("direct batch failed") + + monkeypatch.setattr(spec, "generate", fail) + + with warnings.catch_warnings(), pytest.raises( + RuntimeError, match="direct batch failed" + ): + warnings.simplefilter("ignore", DeprecationWarning) + shell.generate( + cast(torch.LongTensor, input_ids), + do_sample=False, + speculative_draft=spec, + ) + + assert shell._native_generation_engine.spec_strategy is None + + +def test_generation_config_is_normalized_for_selected_dflash( + monkeypatch: pytest.MonkeyPatch, +) -> None: + shell, target = _tiny_moe_shell() + spec = _tiny_speculator(shell, target) + input_ids = torch.tensor([PROMPT, PROMPT], dtype=torch.long) + calls: list[dict[str, Any]] = [] + + def direct_generate(ids: torch.Tensor, **kwargs: Any) -> torch.Tensor: + calls.append({"input_ids": ids, **kwargs}) + return torch.cat([ids, torch.full((2, 4), 7, dtype=ids.dtype)], dim=1) + + monkeypatch.setattr(spec, "generate", direct_generate) + generation_config = SimpleNamespace( + do_sample=False, + temperature=0.0, + top_p=1.0, + top_k=0, + max_new_tokens=4, + eos_token_id=2, + pad_token_id=0, + ) + + with warnings.catch_warnings(): + warnings.simplefilter("ignore", DeprecationWarning) + shell.generate( + cast(torch.LongTensor, input_ids), + generation_config=generation_config, + speculative_draft=spec, + ) + + assert len(calls) == 1 + assert calls[0]["max_new_tokens"] == 4 + assert calls[0]["temperature"] == 0.0 + assert calls[0]["top_k"] == 0 + assert calls[0]["top_p"] == 1.0 + assert calls[0]["stop_token_ids"] == [2] + + +def test_qwen35_sampled_dflash_rejects_before_speculator_execution( + monkeypatch: pytest.MonkeyPatch, +) -> None: + shell = MoE.__new__(MoE) + shell.model = SimpleNamespace( + config=SimpleNamespace(model_type="qwen3_5_moe"), + generate=lambda *_args, **_kwargs: pytest.fail( + "HF fallback was called" + ), + ) + shell.use_native_engine = True + shell._native_generation_engine = SimpleNamespace( + generate=lambda **_kwargs: None + ) + shell.max_seq_length = 64 + resolve_calls: list[object] = [] + + def resolve(value: object) -> None: + resolve_calls.append(value) + + monkeypatch.setattr(shell, "_resolve_spec_strategy", resolve) + + with warnings.catch_warnings(), pytest.raises(ValueError, match="greedy"): + warnings.simplefilter("ignore", DeprecationWarning) + shell.generate( + cast(torch.LongTensor, torch.tensor([[1, 2]])), + do_sample=True, + temperature=0.7, + speculative_draft=object(), + ) + + assert resolve_calls == [] diff --git a/tests/python/dflash/test_request_rng.py b/tests/python/dflash/test_request_rng.py new file mode 100644 index 00000000..9fb7abe3 --- /dev/null +++ b/tests/python/dflash/test_request_rng.py @@ -0,0 +1,299 @@ +from __future__ import annotations + +from collections.abc import Iterable + +import pytest +import torch + +from moe_infinity.spec_decode import DFlashSpeculator, read_dflash_config +from moe_infinity.spec_decode._dflash_sample_ops import acceptance_sampled +from moe_infinity.spec_decode.dflash import SpecSession +from tests.python.dflash.fixtures_tiny import ( + build_tiny_drafter, + build_tiny_target, + make_tiny_drafter_config, +) + +PROMPTS = { + "a": torch.tensor([[3, 7, 11, 2, 5]], dtype=torch.long), + "b": torch.tensor([[13, 1, 9, 4, 6]], dtype=torch.long), +} + + +def _tiny_spec() -> DFlashSpeculator: + target = build_tiny_target(seed=0) + drafter = build_tiny_drafter(target, seed=1) + config = read_dflash_config(make_tiny_drafter_config(target.config)) + return DFlashSpeculator.from_models( + target, drafter, config=config, device="cpu" + ) + + +def _run_to_completion( + spec: DFlashSpeculator, session: SpecSession +) -> list[int]: + while not session.finished: + spec.draft_round(session) + spec.verify_round(session) + return session.output_ids + + +def _run_requests( + request_order: Iterable[str], *, include: Iterable[str] = ("a", "b") +) -> dict[str, list[int]]: + spec = _tiny_spec() + seeds = {"a": 101, "b": 202} + sessions = { + name: spec.begin_session( + PROMPTS[name], + max_new_tokens=18, + temperature=0.8, + top_p=0.9, + generator=torch.Generator().manual_seed(seeds[name]), + ) + for name in include + } + order = list(request_order) + while any(not session.finished for session in sessions.values()): + for name in order: + session = sessions.get(name) + if session is None or session.finished: + continue + spec.draft_round(session) + spec.verify_round(session) + return {name: session.output_ids for name, session in sessions.items()} + + +def test_explicit_session_generator_isolated_from_ambient_rng() -> None: + def generate(ambient_seed: int) -> list[int]: + spec = _tiny_spec() + torch.manual_seed(ambient_seed) + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.8, + top_p=0.9, + generator=torch.Generator().manual_seed(77), + ) + return _run_to_completion(spec, session) + + assert generate(11) == generate(999) + + +def test_explicit_session_generator_does_not_advance_ambient_rng() -> None: + spec = _tiny_spec() + torch.manual_seed(314159) + before = torch.random.get_rng_state().clone() + + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.8, + generator=torch.Generator().manual_seed(77), + ) + _run_to_completion(spec, session) + + assert torch.equal(before, torch.random.get_rng_state()) + + +def test_session_outputs_are_invariant_to_request_order() -> None: + forward = _run_requests(("a", "b")) + reverse = _run_requests(("b", "a")) + + assert forward == reverse + + +def test_session_output_is_invariant_to_unrelated_request_composition() -> None: + alone = _run_requests(("a",), include=("a",))["a"] + composed = _run_requests(("a", "b"))["a"] + + assert alone == composed + + +def test_generate_accepts_an_isolated_request_generator() -> None: + def generate(ambient_seed: int) -> list[int]: + spec = _tiny_spec() + torch.manual_seed(ambient_seed) + return spec.generate( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.8, + top_p=0.9, + generator=torch.Generator().manual_seed(77), + )[0].tolist() + + assert generate(11) == generate(999) + + +def test_session_retains_complete_proposal_probs_until_verification() -> None: + spec = _tiny_spec() + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.8, + generator=torch.Generator().manual_seed(77), + ) + + spec.draft_round(session) + + assert session.pending_draft_probs is not None + assert session.pending_draft_probs.shape == ( + spec.config.block_size - 1, + spec.config.vocab_size, + ) + assert torch.equal( + session.pending_draft_probs.sum(dim=-1), + torch.ones(spec.config.block_size - 1), + ) + + spec.verify_round(session) + assert session.pending_draft_probs is None + + +def test_sampled_verify_refeeds_final_draw_as_next_anchor() -> None: + spec = _tiny_spec() + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.8, + generator=torch.Generator().manual_seed(77), + ) + spec.draft_round(session) + + result = spec.verify_round(session) + + assert not result.finished + assert session.anchor == result.accepted_token_ids[-1] + spec.draft_round(session) + assert session._pending_block is not None + assert int(session._pending_block[0, 0]) == session.anchor + + +def test_greedy_session_does_not_consume_request_generator() -> None: + spec = _tiny_spec() + generator = torch.Generator().manual_seed(77) + before = generator.get_state().clone() + + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.0, + generator=generator, + ) + _run_to_completion(spec, session) + + assert torch.equal(before, generator.get_state()) + + +def test_finished_session_consumes_no_additional_rng() -> None: + spec = _tiny_spec() + generator = torch.Generator().manual_seed(77) + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=8, + temperature=0.8, + generator=generator, + ) + _run_to_completion(spec, session) + after_finish = generator.get_state().clone() + + with pytest.raises(RuntimeError, match="finished session"): + spec.draft_round(session) + + assert torch.equal(after_finish, generator.get_state()) + + +def test_omitted_generator_keeps_global_rng_compatibility() -> None: + def generate() -> list[int]: + spec = _tiny_spec() + torch.manual_seed(77) + session = spec.begin_session( + PROMPTS["a"], max_new_tokens=18, temperature=0.8 + ) + return _run_to_completion(spec, session) + + assert generate() == generate() + + +def test_begin_session_normalizes_negative_top_k_to_disabled() -> None: + spec = _tiny_spec() + + session = spec.begin_session( + PROMPTS["a"], max_new_tokens=8, temperature=0.0, top_k=-1 + ) + + assert session.sampling.top_k == 0 + assert session.sampling.is_greedy + + +def test_acceptance_sampled_rejects_generator_device_mismatch_before_draw() -> ( + None +): + generator = torch.Generator().manual_seed(17) + before = generator.get_state().clone() + draft_probs = torch.empty((1, 2), device="meta") + target_probs = torch.empty((2, 2), device="meta") + drafts = torch.empty((1,), dtype=torch.long, device="meta") + + with pytest.raises( + ValueError, + match="generator device cpu does not match probability device meta", + ): + acceptance_sampled( + draft_probs, target_probs, drafts, generator=generator + ) + + assert torch.equal(before, generator.get_state()) + + +def test_sampled_session_advances_its_request_generator() -> None: + spec = _tiny_spec() + generator = torch.Generator().manual_seed(77) + before = generator.get_state().clone() + + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.8, + generator=generator, + ) + _run_to_completion(spec, session) + + assert not torch.equal(before, generator.get_state()) + + +def test_distinct_request_seeds_produce_non_degenerate_token_streams() -> None: + streams: list[tuple[int, ...]] = [] + states: list[torch.Tensor] = [] + for seed in (7, 17, 27, 37): + spec = _tiny_spec() + generator = torch.Generator().manual_seed(seed) + session = spec.begin_session( + PROMPTS["a"], + max_new_tokens=18, + temperature=0.8, + generator=generator, + ) + streams.append(tuple(_run_to_completion(spec, session))) + states.append(generator.get_state().clone()) + + assert len(set(streams)) > 1 + assert any(not torch.equal(states[0], state) for state in states[1:]) + + +@pytest.mark.skipif(not torch.cuda.is_available(), reason="CUDA is unavailable") +def test_acceptance_sampled_supports_matching_cuda_generator() -> None: + device = torch.device("cuda") + generator = torch.Generator(device=device).manual_seed(17) + draft_probs = torch.tensor([[0.5, 0.5]], device=device) + target_probs = torch.tensor([[0.2, 0.8], [0.7, 0.3]], device=device) + drafts = torch.tensor([0], dtype=torch.long, device=device) + + decisions = [ + acceptance_sampled( + draft_probs, target_probs, drafts, generator=generator + ) + for _ in range(16) + ] + + assert len(decisions) == 16 diff --git a/tests/python/dflash/test_rich_batch_forward.py b/tests/python/dflash/test_rich_batch_forward.py new file mode 100644 index 00000000..59bbd36e --- /dev/null +++ b/tests/python/dflash/test_rich_batch_forward.py @@ -0,0 +1,267 @@ +from __future__ import annotations + +import types + +import pytest +import torch + +from moe_infinity.distributed.expert_executor import DistributedExpertExecutor +from moe_infinity.entrypoints.big_modeling import MoE +from moe_infinity.spec_decode._route_ahead_ctx import route_ahead_context +from moe_infinity.spec_decode.backends_rich import BatchedRichBackend +from moe_infinity.spec_decode.dflash import DFlashSpeculator +from moe_infinity.spec_decode.protocols import ( + RichBatchMetadata, + RichForwardResult, +) +from tests.python.dflash import test_batched_spec as batched + + +class _BatchModel: + def __init__(self) -> None: + self.calls: list[tuple[torch.Tensor, dict[str, object]]] = [] + + def __call__(self, input_ids: torch.Tensor, **kwargs: object) -> object: + self.calls.append((input_ids.clone(), dict(kwargs))) + values = input_ids.to(torch.float32).unsqueeze(-1) + logits = torch.cat((values, values + 100), dim=-1) + return types.SimpleNamespace( + logits=logits, + hidden_states=(values, values + 10), + past_key_values=kwargs.get("past_key_values", "new-cache"), + ) + + def modules(self) -> list[object]: + return [] + + +def _shell(model: _BatchModel) -> MoE: + shell = MoE.__new__(MoE) + shell.model = model + shell._cached_past_key_values = None + shell._native_attention_backend = None + shell._native_mla_cache = None + shell._resolve_native_input_device = lambda: torch.device("cpu") + return shell + + +def test_rich_batch_metadata_rejects_inconsistent_row_layout() -> None: + with pytest.raises(ValueError, match="row_offsets"): + RichBatchMetadata(row_offsets=(0, 2), row_lengths=(2, 1)) + + +def test_native_rich_batch_runs_one_forward_and_preserves_rows_masks_positions() -> ( + None +): + model = _BatchModel() + shell = _shell(model) + input_ids = torch.tensor([[0, 4, 5], [7, 8, 9]]) + mask = torch.tensor([[0, 1, 1], [1, 1, 1]]) + positions = torch.tensor([[0, 0, 1], [3, 4, 5]]) + cache = object() + metadata = RichBatchMetadata( + row_offsets=(0, 2, 5), + row_lengths=(2, 3), + attention_mask=mask, + position_ids=positions, + cache_handles=(cache, cache), + request_contexts=("left", "right"), + route_contexts=("route-left", "route-right"), + ) + + result = shell._native_model_forward_rich(input_ids, metadata) + + assert isinstance(result, RichForwardResult) + assert len(model.calls) == 1 + called_ids, kwargs = model.calls[0] + assert torch.equal(called_ids, input_ids) + assert torch.equal(kwargs["attention_mask"], mask) + assert torch.equal(kwargs["position_ids"], positions) + assert result.logits.shape == (2, 3, 2) + assert result.logits[0, 1].tolist() == [4.0, 104.0] + assert result.logits[1, 0].tolist() == [7.0, 107.0] + assert result.hidden_states[1][0, 1, 0].item() == 14 + assert result.cache_handles == (result.cache_handle, result.cache_handle) + assert result.row_offsets == (0, 2, 5) + assert result.row_lengths == (2, 3) + + +def test_native_rich_legacy_list_keeps_singleton_tuple_contract() -> None: + model = _BatchModel() + shell = _shell(model) + + result = shell._native_model_forward_rich([4, 5]) + + assert isinstance(result, tuple) + assert len(result) == 3 + assert result[0].shape == (1, 2, 2) + + +def _tiny_rich_spec( + *, supports_batch: bool = True +) -> tuple[DFlashSpeculator, object, list[int]]: + base, target = batched._tiny_spec() + shell = MoE.__new__(MoE) + shell.model = target + shell._cached_past_key_values = None + shell._native_attention_backend = None + shell._native_mla_cache = None + shell._resolve_native_input_device = lambda: torch.device("cpu") + shell._configure_hook = lambda _ids: None + calls = [0] + original = shell._native_model_forward_rich + + def counted(*args: object, **kwargs: object) -> object: + calls[0] += 1 + return original(*args, **kwargs) + + shell._native_model_forward_rich = counted + shell._supports_native_rich_batch = lambda: supports_batch + spec = DFlashSpeculator.from_models( + shell, base.draft, config=base.config, device="cpu" + ) + return spec, target, calls + + +def test_rich_backend_is_physical_only_for_declared_row_aware_wrapper() -> None: + spec, _, _ = _tiny_rich_spec() + backend = BatchedRichBackend(spec) + + assert backend.capabilities.supports_batch + assert backend.capabilities.supports_rich_forward + assert backend.wrapper_supported + + spec.moe._supports_native_rich_batch = lambda: False + unsupported = BatchedRichBackend(spec) + assert not unsupported.wrapper_supported + + +def test_rich_generate_uses_real_multirow_target_forwards_and_matches_independent() -> ( + None +): + spec, target, calls = _tiny_rich_spec() + prompts = [batched.PROMPT_A, batched.PROMPT_B] + input_ids, mask, width = batched._left_pad(prompts) + model_batch_sizes: list[int] = [] + hook = target.register_forward_pre_hook( + lambda _module, args: model_batch_sizes.append(int(args[0].shape[0])) + ) + + try: + output = spec.generate( + input_ids, max_new_tokens=[5, 7], attention_mask=mask + ) + finally: + hook.remove() + rows = batched._batched_new_tokens(output, spec, width) + + assert spec.rich_forward_batched is True + assert calls[0] == len(model_batch_sizes) > 0 + assert all(size > 1 for size in model_batch_sizes) + assert all(size == 2 for size in spec.rich_forward_batch_sizes) + + expected = [] + for prompt, budget in zip(prompts, (5, 7)): + single, _, _ = _tiny_rich_spec() + result = single.generate(torch.tensor([prompt]), max_new_tokens=budget) + expected.append(result[0, len(prompt) :].tolist()) + assert rows == expected + + +def _run_named_rich_rows(names: list[str]) -> dict[str, list[int]]: + prompts = { + "a": [3, 7, 11, 2, 5], + "b": [41, 6], + "c": [10, 20, 30, 40], + } + spec, target, rich_calls = _tiny_rich_spec() + input_ids, mask, width = batched._left_pad( + [prompts[name] for name in names] + ) + model_batch_sizes: list[int] = [] + hook = target.register_forward_pre_hook( + lambda _module, args: model_batch_sizes.append(int(args[0].shape[0])) + ) + + try: + output = spec.generate(input_ids, max_new_tokens=6, attention_mask=mask) + finally: + hook.remove() + rows = batched._batched_new_tokens(output, spec, width) + + assert spec.rich_forward_batched is True + assert ( + rich_calls[0] + == len(model_batch_sizes) + == len(spec.rich_forward_batch_sizes) + ) + assert rich_calls[0] >= 1 + assert all(size > 1 for size in model_batch_sizes) + assert all( + size == len(names) and size > 1 + for size in spec.rich_forward_batch_sizes + ) + assert len({tuple(row) for row in rows}) == len(rows) + return {name: rows[row] for row, name in enumerate(names)} + + +def test_rich_physical_rows_match_independent_and_are_order_composition_invariant() -> ( + None +): + forward = _run_named_rich_rows(["a", "b", "c"]) + reverse = _run_named_rich_rows(["c", "b", "a"]) + composed = _run_named_rich_rows(["a", "c"]) + + independent: dict[str, list[int]] = {} + prompts = {"a": [3, 7, 11, 2, 5], "b": [41, 6], "c": [10, 20, 30, 40]} + for name, prompt in prompts.items(): + spec, _, _ = _tiny_rich_spec() + output = spec.generate(torch.tensor([prompt]), max_new_tokens=6) + independent[name] = output[0, len(prompt) :].tolist() + + assert forward == reverse == independent + assert composed == {"a": independent["a"], "c": independent["c"]} + + +@pytest.mark.parametrize("wrapper_kind", ["mla", "hybrid"]) +def test_unsupported_mla_or_hybrid_wrapper_falls_back_per_request( + wrapper_kind: str, + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec, _, calls = _tiny_rich_spec() + if wrapper_kind == "mla": + spec.moe._get_mla_attention_modules = lambda: [object()] + else: + monkeypatch.setattr( + spec.moe.model.config, "hybrid_attention", True, raising=False + ) + spec.moe._supports_native_rich_batch = types.MethodType( + MoE._supports_native_rich_batch, spec.moe + ) + assert spec.moe._supports_native_rich_batch() is False + ids, mask, _ = batched._left_pad([batched.PROMPT_A, batched.PROMPT_B]) + + spec.generate(ids, max_new_tokens=2, attention_mask=mask) + + assert spec.rich_forward_batched is False + assert spec.rich_forward_batch_sizes == [] + assert calls[0] >= 2 + + +def test_executor_evidence_retains_every_rich_row_layer_union() -> None: + executor = DistributedExpertExecutor.__new__(DistributedExpertExecutor) + executor.prefetcher = None + mask = torch.tensor( + [[1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 0, 0]], dtype=torch.bool + ) + + with route_ahead_context(row_offsets=(0, 2, 4)): + executor._maybe_route_ahead_prefetch(7, mask, 3) + + assert executor.last_executor_evidence.actual_expert_union == frozenset( + {(7, 0), (7, 1), (7, 2)} + ) + assert ( + executor.last_executor_evidence.actual_expert_union_by_row + == frozenset({(0, 7, 0), (0, 7, 1), (1, 7, 0), (1, 7, 2)}) + ) diff --git a/tests/python/dflash/test_session_driver.py b/tests/python/dflash/test_session_driver.py new file mode 100644 index 00000000..b373a66e --- /dev/null +++ b/tests/python/dflash/test_session_driver.py @@ -0,0 +1,547 @@ +from __future__ import annotations + +from dataclasses import dataclass, field +from typing import Callable, Hashable + +import pytest +import torch + +import moe_infinity.spec_decode.session_driver as driver_module +from moe_infinity.spec_decode import DFlashSpeculator, read_dflash_config +from moe_infinity.spec_decode.backends import ( + DFlashExecutionBackend, + ExecutionBackend, +) +from moe_infinity.spec_decode.dflash import VerifyResult +from moe_infinity.spec_decode.protocols import ( + BackendCapabilities, + RequestSpec, + SamplingContext, + SessionRoundResult, + SessionTrace, +) +from moe_infinity.spec_decode.session_driver import ( + SessionDriver, + UnsupportedRequestError, +) +from tests.python.dflash.fixtures_tiny import ( + build_tiny_drafter, + build_tiny_target, + make_tiny_drafter_config, +) + + +@dataclass +class _State: + request: RequestSpec + output: list[int] = field(default_factory=list) + pending: bool = False + finished: bool = False + released: bool = False + restored: bool = False + rounds: int = 0 + + +class _FakeBackend: + def __init__( + self, + name: str, + events: list[tuple[str, str, str]], + *, + supports_sampling: bool = True, + accepts: Callable[[RequestSpec], bool] = lambda _request: True, + cohort: Callable[[RequestSpec], Hashable] = lambda _request: "default", + fail_draft_for: frozenset[str] = frozenset(), + draft_error: BaseException | None = None, + no_progress_for: frozenset[str] = frozenset(), + fail_restore_for: frozenset[str] = frozenset(), + fail_release_for: frozenset[str] = frozenset(), + ) -> None: + self.name = name + self.events = events + self.accepts = accepts + self.cohort = cohort + self.fail_draft_for = fail_draft_for + self.draft_error = draft_error + self.no_progress_for = no_progress_for + self.fail_restore_for = fail_restore_for + self.fail_release_for = fail_release_for + self.sessions: list[_State] = [] + self.capabilities = BackendCapabilities( + supports_batch=False, + supports_sampling=supports_sampling, + supports_ragged_rows=True, + cache_kind="dense_dynamic", + supports_route_ahead=False, + supports_rich_forward=True, + ) + + def supports(self, request: RequestSpec) -> bool: + self.events.append(("supports", self.name, request.request_id)) + return self.accepts(request) + + def cohort_key(self, request: RequestSpec) -> Hashable: + return self.cohort(request) + + def prefill(self, request: RequestSpec) -> _State: + self.events.append(("prefill", self.name, request.request_id)) + state = _State(request=request, finished=request.max_new_tokens == 0) + self.sessions.append(state) + return state + + def draft(self, session: _State) -> None: + assert not session.restored, "driver resumed an abort-restored session" + self.events.append(("draft", self.name, session.request.request_id)) + if session.finished: + raise AssertionError("finished row was drafted") + session.pending = True + if session.request.request_id in self.fail_draft_for: + # Deliberately corrupt tentative state. The driver must restore it. + session.output.append(999) + raise self.draft_error or RuntimeError("draft failed") + + def verify(self, session: _State) -> SessionRoundResult: + assert not session.restored, "driver resumed an abort-restored session" + self.events.append(("verify", self.name, session.request.request_id)) + assert session.pending + if session.request.request_id in self.no_progress_for: + session.pending = False + return SessionRoundResult( + accepted_draft_count=0, + committed_token_ids=(777,), + next_anchor=777, + target_cache_length=len(session.request.prompt_token_ids), + emitted_length=len(session.output), + finished=False, + finish_reason=None, + ) + token = 10 + session.rounds + session.output.append(token) + session.rounds += 1 + session.pending = False + stop = token in session.request.stop_token_ids + session.finished = ( + stop or len(session.output) >= session.request.max_new_tokens + ) + return SessionRoundResult( + accepted_draft_count=0, + committed_token_ids=(token,), + next_anchor=token, + target_cache_length=len(session.request.prompt_token_ids) + + len(session.output), + emitted_length=len(session.output), + finished=session.finished, + finish_reason=("stop" if stop else "length") + if session.finished + else None, + ) + + def snapshot(self, session: _State) -> tuple[list[int], bool, bool, int]: + return ( + list(session.output), + session.pending, + session.finished, + session.rounds, + ) + + def restore( + self, + session: _State, + snapshot: tuple[list[int], bool, bool, int], + ) -> None: + assert not session.restored, "abort-only restore called more than once" + self.events.append(("restore", self.name, session.request.request_id)) + output, pending, finished, rounds = snapshot + session.output = output + session.pending = pending + session.finished = finished + session.rounds = rounds + session.restored = True + if session.request.request_id in self.fail_restore_for: + raise RuntimeError(f"restore failed: {session.request.request_id}") + + def is_finished(self, session: _State) -> bool: + assert not session.restored, "driver resumed an abort-restored session" + return session.finished + + def output_token_ids(self, session: _State) -> tuple[int, ...]: + assert not session.restored, "driver resumed an abort-restored session" + return tuple(session.output) + + def trace(self, session: _State) -> SessionTrace: + assert not session.restored, "driver resumed an abort-restored session" + return SessionTrace( + request_id=session.request.request_id, + backend=self.name, + cache_kind=self.capabilities.cache_kind, + sampled=session.request.is_sampled, + round_count=session.rounds, + emitted=len(session.output), + finish_reason=( + "stop" + if session.output + and session.output[-1] in session.request.stop_token_ids + else "length" + ), + ) + + def release(self, session: _State) -> None: + self.events.append(("release", self.name, session.request.request_id)) + session.released = True + if session.request.request_id in self.fail_release_for: + raise RuntimeError(f"release failed: {session.request.request_id}") + + +def _request( + request_id: str, + *, + budget: int = 2, + sampling: SamplingContext | None = None, + stops: frozenset[int] = frozenset(), +) -> RequestSpec: + return RequestSpec( + request_id=request_id, + prompt_token_ids=(1, 2, 3), + max_new_tokens=budget, + stop_token_ids=stops, + sampling=sampling or SamplingContext(), + ) + + +def test_driver_creates_one_session_per_row() -> None: + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend("native", events) + requests = [_request("r0", budget=1), _request("r1", budget=2)] + + results = SessionDriver([backend]).run(requests) + + assert isinstance(backend, ExecutionBackend) + assert [state.request for state in backend.sessions] == requests + assert [result.request_id for result in results] == ["r0", "r1"] + assert [result.output_token_ids for result in results] == [(10,), (10, 11)] + assert all(state.released for state in backend.sessions) + + +def test_backend_selected_before_prefill() -> None: + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend("native", events) + + SessionDriver([backend]).run([_request("r0"), _request("r1")]) + + first_prefill = next( + i for i, event in enumerate(events) if event[0] == "prefill" + ) + assert [event[0] for event in events[:first_prefill]] == [ + "supports", + "supports", + ] + + +def test_incompatible_rows_split() -> None: + events: list[tuple[str, str, str]] = [] + even = _FakeBackend( + "even", events, accepts=lambda request: request.request_id.endswith("0") + ) + odd = _FakeBackend( + "odd", events, accepts=lambda request: request.request_id.endswith("1") + ) + driver = SessionDriver([even, odd]) + + results = driver.run([_request("r0", budget=1), _request("r1", budget=1)]) + + assert [result.backend for result in results] == ["even", "odd"] + assert [ + (cohort.backend, cohort.row_indices) for cohort in driver.last_cohorts + ] == [ + ("even", (0,)), + ("odd", (1,)), + ] + + +def test_unsupported_sampling_no_greedy_downgrade() -> None: + events: list[tuple[str, str, str]] = [] + greedy = _FakeBackend("greedy", events, supports_sampling=False) + sampled = _FakeBackend("sampled", events, supports_sampling=True) + sampling = SamplingContext( + temperature=0.8, + top_p=0.9, + generator=torch.Generator().manual_seed(17), + ) + + result = SessionDriver([greedy, sampled]).run( + [_request("sampled-row", budget=1, sampling=sampling)] + )[0] + + assert result.backend == "sampled" + assert sampled.sessions[0].request.sampling is sampling + assert greedy.sessions == [] + + +def test_unsupported_sampling_fails_before_output() -> None: + events: list[tuple[str, str, str]] = [] + greedy = _FakeBackend("greedy", events, supports_sampling=False) + sampled = SamplingContext(temperature=0.7) + + with pytest.raises(UnsupportedRequestError, match="sampled-row"): + SessionDriver([greedy]).run( + [_request("greedy-row"), _request("sampled-row", sampling=sampled)] + ) + + assert greedy.sessions == [] + assert all(event[0] == "supports" for event in events) + + +def test_failure_before_verify_no_partial_output() -> None: + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend( + "native", events, fail_draft_for=frozenset({"broken"}) + ) + driver = SessionDriver([backend]) + + with pytest.raises(RuntimeError, match="draft failed"): + driver.run([_request("healthy"), _request("broken")]) + + assert driver.last_results == () + assert [state.output for state in backend.sessions] == [[], []] + assert all(state.released for state in backend.sessions) + assert not any( + event[0] == "verify" and event[2] == "broken" for event in events + ) + + +def test_snapshot_restore_contract_is_abort_only() -> None: + assert "abort-only" in (ExecutionBackend.snapshot.__doc__ or "") + assert "must not resume" in (ExecutionBackend.restore.__doc__ or "") + + +def test_driver_never_resumes_or_restores_a_failed_session_twice() -> None: + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend( + "native", events, fail_draft_for=frozenset({"broken"}) + ) + + with pytest.raises(RuntimeError, match="draft failed"): + SessionDriver([backend]).run([_request("broken")]) + + assert [event[0] for event in events].count("restore") == 1 + restore_index = next( + i for i, event in enumerate(events) if event[0] == "restore" + ) + assert all(event[0] == "release" for event in events[restore_index + 1 :]) + + +def test_no_progress_round_aborts_all_sessions_and_publishes_no_results() -> ( + None +): + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend( + "native", events, no_progress_for=frozenset({"stuck"}) + ) + driver = SessionDriver([backend]) + + with pytest.raises( + driver_module.BackendProgressError, + match="backend native made no progress for request 'stuck'", + ): + driver.run([_request("stuck"), _request("other")]) + + assert driver.last_results == () + assert all(state.restored and state.released for state in backend.sessions) + assert [state.output for state in backend.sessions] == [[], []] + + +def test_primary_execution_error_survives_all_cleanup_failures() -> None: + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend( + "native", + events, + fail_draft_for=frozenset({"broken"}), + fail_restore_for=frozenset({"healthy"}), + fail_release_for=frozenset({"healthy"}), + ) + driver = SessionDriver([backend]) + + with pytest.raises(RuntimeError) as caught: + driver.run([_request("healthy"), _request("broken")]) + + assert str(caught.value) == "draft failed" + assert driver.last_results == () + assert [event[2] for event in events if event[0] == "restore"] == [ + "healthy", + "broken", + ] + assert [event[2] for event in events if event[0] == "release"] == [ + "healthy", + "broken", + ] + assert any( + "restore failed: healthy" in note for note in caught.value.__notes__ + ) + assert any( + "release failed: healthy" in note for note in caught.value.__notes__ + ) + + +def test_primary_error_survives_cleanup_when_add_note_is_unavailable() -> None: + class LegacyRuntimeError(RuntimeError): + def __getattribute__(self, name: str) -> object: + if name == "add_note": + return None + return super().__getattribute__(name) + + events: list[tuple[str, str, str]] = [] + primary = LegacyRuntimeError("legacy primary") + backend = _FakeBackend( + "native", + events, + fail_draft_for=frozenset({"broken"}), + draft_error=primary, + fail_release_for=frozenset({"broken"}), + ) + + with pytest.raises(LegacyRuntimeError, match="legacy primary") as caught: + SessionDriver([backend]).run([_request("broken")]) + + assert caught.value is primary + assert backend.sessions[0].released + cleanup_context = caught.value.__context__ + assert isinstance(cleanup_context, driver_module.SessionCleanupError) + assert [str(error) for error in cleanup_context.errors] == [ + "release failed: broken" + ] + assert "release failed: broken" in str(cleanup_context) + + +def test_successful_execution_surfaces_cleanup_errors_after_all_releases() -> ( + None +): + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend( + "native", + events, + fail_release_for=frozenset({"r0", "r1"}), + ) + driver = SessionDriver([backend]) + + with pytest.raises(driver_module.SessionCleanupError) as caught: + driver.run([_request("r0", budget=1), _request("r1", budget=1)]) + + assert driver.last_results == () + assert [event[2] for event in events if event[0] == "release"] == [ + "r0", + "r1", + ] + assert [str(error) for error in caught.value.errors] == [ + "release failed: r0", + "release failed: r1", + ] + + +def test_finished_rows_inactive() -> None: + events: list[tuple[str, str, str]] = [] + backend = _FakeBackend("native", events) + + results = SessionDriver([backend]).run( + [_request("short", budget=1), _request("long", budget=3)] + ) + + drafted = [event[2] for event in events if event[0] == "draft"] + assert drafted.count("short") == 1 + assert drafted.count("long") == 3 + assert [len(result.output_token_ids) for result in results] == [1, 3] + + +def test_dflash_backend_is_a_per_request_session_adapter() -> None: + target = build_tiny_target(seed=0) + drafter = build_tiny_drafter(target, seed=1) + config = read_dflash_config(make_tiny_drafter_config(target.config)) + speculator = DFlashSpeculator.from_models( + target, drafter, config=config, device="cpu" + ) + backend = DFlashExecutionBackend(speculator) + + results = SessionDriver([backend]).run( + [_request("r0", budget=3), _request("r1", budget=5)] + ) + + assert [len(result.output_token_ids) for result in results] == [3, 5] + assert [result.request_id for result in results] == ["r0", "r1"] + assert all(result.backend == "dflash-per-request" for result in results) + assert all( + result.trace.round_count == len(result.rounds) for result in results + ) + + +def test_dflash_backend_empty_verify_is_a_finished_noop( + monkeypatch: pytest.MonkeyPatch, +) -> None: + target = build_tiny_target(seed=0) + drafter = build_tiny_drafter(target, seed=1) + config = read_dflash_config(make_tiny_drafter_config(target.config)) + speculator = DFlashSpeculator.from_models( + target, drafter, config=config, device="cpu" + ) + backend = DFlashExecutionBackend(speculator) + session = backend.prefill(_request("empty", budget=0)) + monkeypatch.setattr( + speculator, + "verify_round", + lambda _session: VerifyResult( + accepted_token_ids=[], + accept=0, + verified_accept=5, + committed_count=0, + finished=True, + ), + ) + + result = backend.verify(session) + + assert result.accepted_draft_count == 5 + assert result.committed_token_ids == () + assert result.cached_token_count == 0 + assert result.emitted_length == 0 + assert result.next_anchor is None + assert result.finished + + +def test_spec_session_clear_pending_encapsulates_abort_state() -> None: + target = build_tiny_target(seed=0) + drafter = build_tiny_drafter(target, seed=1) + config = read_dflash_config(make_tiny_drafter_config(target.config)) + speculator = DFlashSpeculator.from_models( + target, drafter, config=config, device="cpu" + ) + session = speculator.begin_session( + torch.tensor([[1, 2, 3]], dtype=torch.long), max_new_tokens=4 + ) + speculator.draft_round(session) + + assert session.has_pending_draft + session.clear_pending() + + assert not session.has_pending_draft + with pytest.raises(RuntimeError, match="without a pending draft"): + speculator.verify_round(session) + + +def test_verify_round_clamps_over_budget_defensive_entry() -> None: + target = build_tiny_target(seed=0) + drafter = build_tiny_drafter(target, seed=1) + config = read_dflash_config(make_tiny_drafter_config(target.config)) + speculator = DFlashSpeculator.from_models( + target, drafter, config=config, device="cpu" + ) + session = speculator.begin_session( + torch.tensor([[1, 2, 3]], dtype=torch.long), max_new_tokens=1 + ) + speculator.draft_round(session) + output_before = session.output_ids + session.emitted.extend([55, 56]) + + result = speculator.verify_round(session) + + assert result.accepted_token_ids == [] + assert result.committed_count == 0 + assert result.finished + assert session.output_ids == output_before diff --git a/tests/python/dflash/test_single_session_parity.py b/tests/python/dflash/test_single_session_parity.py new file mode 100644 index 00000000..c3eb4aef --- /dev/null +++ b/tests/python/dflash/test_single_session_parity.py @@ -0,0 +1,512 @@ +from __future__ import annotations + +from collections.abc import Callable +from types import SimpleNamespace +from typing import Any + +import pytest +import torch + +import moe_infinity.spec_decode.dflash as dflash_module +from moe_infinity.spec_decode import ( + DFlashExecutionBackend, + DFlashSpeculator, + RequestSpec, + read_dflash_config, +) +from tests.python.dflash.fixtures_tiny import ( + TINY_BLOCK_SIZE, + TINY_HIDDEN, + TINY_VOCAB, + build_tiny_drafter, + build_tiny_target, + make_tiny_drafter_config, + plain_greedy_decode, +) + +PROMPT = torch.tensor([[3, 7, 11, 2, 5]], dtype=torch.long) +PROMPT_LEN = int(PROMPT.shape[1]) +EOS_ID = 62 + + +def _tiny_spec(seed: int = 0) -> DFlashSpeculator: + target = build_tiny_target(seed=seed) + drafter = build_tiny_drafter(target, seed=seed + 1) + config = read_dflash_config(make_tiny_drafter_config(target.config)) + return DFlashSpeculator.from_models( + target, drafter, config=config, device="cpu" + ) + + +def _observe_session_driver( + monkeypatch: pytest.MonkeyPatch, spec: DFlashSpeculator +) -> dict[str, int]: + calls = {"begin": 0, "draft": 0, "verify": 0} + real_begin = spec.begin_session + real_draft = spec.draft_round + real_verify = spec.verify_round + + def begin(*args: Any, **kwargs: Any): + calls["begin"] += 1 + return real_begin(*args, **kwargs) + + def draft(*args: Any, **kwargs: Any): + calls["draft"] += 1 + return real_draft(*args, **kwargs) + + def verify(*args: Any, **kwargs: Any): + calls["verify"] += 1 + return real_verify(*args, **kwargs) + + monkeypatch.setattr(spec, "begin_session", begin) + monkeypatch.setattr(spec, "draft_round", draft) + monkeypatch.setattr(spec, "verify_round", verify) + return calls + + +def _run_direct_session( + spec: DFlashSpeculator, + *, + max_new_tokens: int, + temperature: float = 0.0, + top_k: int = 0, + top_p: float = 1.0, + generator: torch.Generator | None = None, +) -> list[int]: + session = spec.begin_session( + PROMPT.clone(), + max_new_tokens=max_new_tokens, + temperature=temperature, + top_k=top_k, + top_p=top_p, + generator=generator, + ) + while not session.finished and len(session.output_ids) < max_new_tokens: + spec.draft_round(session) + spec.verify_round(session) + return session.output_ids + + +def _cache_length(cache: Any) -> int | None: + return None if cache is None else int(cache.get_seq_length()) + + +def test_dense_cache_length_mismatch_raises_runtime_error( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = _tiny_spec() + original_rollback = dflash_module.rollback_target_cache + + def truncate_too_far(cache: Any, snapshot: Any, **kwargs: Any) -> None: + original_rollback(cache, snapshot, **kwargs) + cache.crop(int(kwargs["prev_start"])) + + monkeypatch.setattr( + dflash_module, "rollback_target_cache", truncate_too_far + ) + session = spec.begin_session(PROMPT.clone(), max_new_tokens=2) + spec.draft_round(session) + + with pytest.raises(RuntimeError, match="target cache length invariant"): + spec.verify_round(session) + + +class _ScriptedHead: + def __init__(self, draft_fn: Callable[[int], list[int]]) -> None: + self.draft_fn = draft_fn + self.calls = 0 + + def __call__(self, hidden: torch.Tensor) -> torch.Tensor: + drafts = [int(token) for token in self.draft_fn(self.calls)] + self.calls += 1 + assert len(drafts) == TINY_BLOCK_SIZE - 1 + logits = torch.zeros( + 1, + hidden.shape[1], + TINY_VOCAB, + dtype=hidden.dtype, + device=hidden.device, + ) + offset = hidden.shape[1] - (TINY_BLOCK_SIZE - 1) + for index, token in enumerate(drafts): + logits[0, offset + index, token] = 1.0 + return logits + + +def _install_scripted_drafter( + monkeypatch: pytest.MonkeyPatch, + spec: DFlashSpeculator, + draft_fn: Callable[[int], list[int]], +) -> None: + monkeypatch.setattr( + spec, + "_run_drafter", + lambda block, context_feature, start, draft_kv: torch.zeros( + 1, TINY_BLOCK_SIZE, TINY_HIDDEN + ), + ) + monkeypatch.setattr(spec, "lm_head", _ScriptedHead(draft_fn)) + + +def _force_target_argmax( + monkeypatch: pytest.MonkeyPatch, + spec: DFlashSpeculator, + *, + prefill: dict[int, int] | None = None, + verify: dict[int, int] | None = None, +) -> None: + original = spec._forward_target + + def wrapped(input_ids, past_key_values=None, logits_to_keep=0): + logits, hidden, cache = original( + input_ids, + past_key_values=past_key_values, + logits_to_keep=logits_to_keep, + ) + overrides = prefill if int(logits_to_keep) == 1 else verify + if overrides: + logits = logits.clone() + for row, token in overrides.items(): + logits[0, row, :] = -1e9 + logits[0, row, int(token)] = 1e9 + return logits, hidden, cache + + monkeypatch.setattr(spec, "_forward_target", wrapped) + + +def test_generate_single_delegates_through_real_session_rounds( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = _tiny_spec() + calls = _observe_session_driver(monkeypatch, spec) + + output = spec.generate(PROMPT.clone(), max_new_tokens=14) + + assert calls["begin"] == 1 + assert calls["draft"] == calls["verify"] == len(spec.step_trace) + assert calls["verify"] > 0 + assert output.shape == (1, PROMPT_LEN + 14) + + +@pytest.mark.parametrize("budget", [0, 1]) +def test_zero_and_one_token_budgets_prefill_without_unverified_rounds( + monkeypatch: pytest.MonkeyPatch, budget: int +) -> None: + spec = _tiny_spec() + calls = _observe_session_driver(monkeypatch, spec) + + output = spec.generate(PROMPT.clone(), max_new_tokens=budget) + + assert calls == {"begin": 1, "draft": 0, "verify": 0} + assert torch.equal(output[:, :PROMPT_LEN], PROMPT) + assert output.shape == (1, PROMPT_LEN + budget) + assert output.dtype == PROMPT.dtype + assert output.device == PROMPT.device + assert spec.step_trace == [] + assert spec.last_generated_lengths is None + assert _cache_length(spec.last_target_cache) == PROMPT_LEN + assert spec.last_draft_cache is None + + +def test_greedy_generate_matches_direct_session_diagnostics_exactly() -> None: + budget = 17 + generated_spec = _tiny_spec() + generated_stats = generated_spec.enable_route_ahead_stats() + generated = generated_spec.generate(PROMPT.clone(), max_new_tokens=budget) + + session_spec = _tiny_spec() + session_stats = session_spec.enable_route_ahead_stats() + session_ids = _run_direct_session( + session_spec, max_new_tokens=budget, temperature=0.0 + ) + + assert generated[0, PROMPT_LEN:].tolist() == session_ids + assert generated_spec.step_trace == session_spec.step_trace + assert _cache_length(generated_spec.last_target_cache) == _cache_length( + session_spec.last_target_cache + ) + assert _cache_length(generated_spec.last_draft_cache) == _cache_length( + session_spec.last_draft_cache + ) + assert generated_stats.as_dict() == session_stats.as_dict() + assert generated_spec.last_generated_lengths is None + assert generated.shape == (1, PROMPT_LEN + budget) + assert generated.dtype == PROMPT.dtype + assert generated.device == PROMPT.device + + +def test_sampled_generate_matches_direct_session_with_explicit_generator() -> ( + None +): + budget = 18 + generate_rng = torch.Generator().manual_seed(917) + session_rng = torch.Generator().manual_seed(917) + + generated_spec = _tiny_spec() + generated = generated_spec.generate( + PROMPT.clone(), + max_new_tokens=budget, + temperature=0.8, + top_k=-1, + top_p=0.9, + generator=generate_rng, + ) + session_spec = _tiny_spec() + session_ids = _run_direct_session( + session_spec, + max_new_tokens=budget, + temperature=0.8, + top_k=-1, + top_p=0.9, + generator=session_rng, + ) + + assert generated[0, PROMPT_LEN:].tolist() == session_ids + assert torch.equal(generate_rng.get_state(), session_rng.get_state()) + assert generated_spec.step_trace == session_spec.step_trace + assert _cache_length(generated_spec.last_target_cache) == _cache_length( + session_spec.last_target_cache + ) + assert generated_spec.last_generated_lengths is None + + +def test_immediate_eos_delegates_to_begin_without_a_round( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = _tiny_spec() + _force_target_argmax(monkeypatch, spec, prefill={-1: EOS_ID}) + calls = _observe_session_driver(monkeypatch, spec) + + output = spec.generate( + PROMPT.clone(), max_new_tokens=16, stop_token_ids=[EOS_ID] + ) + + assert output[0, PROMPT_LEN:].tolist() == [EOS_ID] + assert calls == {"begin": 1, "draft": 0, "verify": 0} + assert spec.step_trace == [] + assert _cache_length(spec.last_target_cache) == PROMPT_LEN + + +@pytest.mark.parametrize("eos_kind", ["accepted_draft", "bonus"]) +def test_eos_inside_verified_commit_preserves_cache_and_trace( + monkeypatch: pytest.MonkeyPatch, eos_kind: str +) -> None: + spec = _tiny_spec() + target = spec.target + greedy = plain_greedy_decode(target, PROMPT, max_new_tokens=16)[0].tolist() + if eos_kind == "accepted_draft": + other = 40 + drafts = [ + greedy[PROMPT_LEN + 1], + greedy[PROMPT_LEN + 2], + EOS_ID, + (other + 1) % TINY_VOCAB, + ] + [0] * (TINY_BLOCK_SIZE - 5) + verify = {2: EOS_ID, 3: other} + expected_cache_length = PROMPT_LEN + 4 + else: + drafts = greedy[PROMPT_LEN + 1 : PROMPT_LEN + TINY_BLOCK_SIZE] + verify = {2: EOS_ID} + expected_cache_length = PROMPT_LEN + 3 + _install_scripted_drafter(monkeypatch, spec, lambda _step: drafts) + _force_target_argmax(monkeypatch, spec, verify=verify) + calls = _observe_session_driver(monkeypatch, spec) + + output = spec.generate( + PROMPT.clone(), max_new_tokens=40, stop_token_ids=[EOS_ID] + ) + + assert output[0, -1].item() == EOS_ID + assert calls == {"begin": 1, "draft": 1, "verify": 1} + assert len(spec.step_trace) == 1 + assert _cache_length(spec.last_target_cache) == expected_cache_length + trace = spec.step_trace[0] + assert trace.start == trace.prev_start + trace.accept + 1 + assert trace.target_cache_len == trace.start + + +def test_backend_reports_true_acceptance_for_eos_truncated_round( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = _tiny_spec() + greedy = plain_greedy_decode(spec.target, PROMPT, max_new_tokens=16)[ + 0 + ].tolist() + drafts = [ + greedy[PROMPT_LEN + 1], + greedy[PROMPT_LEN + 2], + EOS_ID, + ] + [0] * (TINY_BLOCK_SIZE - 4) + _install_scripted_drafter(monkeypatch, spec, lambda _step: drafts) + _force_target_argmax(monkeypatch, spec, verify={2: EOS_ID, 3: 40}) + backend = DFlashExecutionBackend(spec) + session = backend.prefill( + RequestSpec( + request_id="eos-truncated", + prompt_token_ids=tuple(PROMPT[0].tolist()), + max_new_tokens=40, + stop_token_ids=frozenset({EOS_ID}), + ) + ) + + backend.draft(session) + result = backend.verify(session) + + assert result.committed_token_ids[-1] == EOS_ID + assert result.accepted_draft_count == 3 + assert len(result.committed_token_ids) == 3 + assert result.next_anchor is None + assert result.finished + + +def test_backend_reports_true_acceptance_for_budget_truncated_round( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = _tiny_spec() + greedy = plain_greedy_decode(spec.target, PROMPT, max_new_tokens=16)[ + 0 + ].tolist() + drafts = greedy[PROMPT_LEN + 1 : PROMPT_LEN + TINY_BLOCK_SIZE] + _install_scripted_drafter(monkeypatch, spec, lambda _step: drafts) + backend = DFlashExecutionBackend(spec) + session = backend.prefill( + RequestSpec( + request_id="budget-truncated", + prompt_token_ids=tuple(PROMPT[0].tolist()), + max_new_tokens=2, + ) + ) + + backend.draft(session) + result = backend.verify(session) + + assert result.accepted_draft_count == TINY_BLOCK_SIZE - 1 + assert len(result.committed_token_ids) == 1 + assert result.next_anchor is None + assert result.finished + + +def test_partial_final_block_delegates_and_crops_to_budget( + monkeypatch: pytest.MonkeyPatch, +) -> None: + spec = _tiny_spec() + target = spec.target + budget = TINY_BLOCK_SIZE + 5 + greedy = plain_greedy_decode(target, PROMPT, max_new_tokens=25)[0].tolist() + + def drafts(step: int) -> list[int]: + anchor_index = PROMPT_LEN + TINY_BLOCK_SIZE * step + return greedy[anchor_index + 1 : anchor_index + TINY_BLOCK_SIZE] + + _install_scripted_drafter(monkeypatch, spec, drafts) + calls = _observe_session_driver(monkeypatch, spec) + + output = spec.generate(PROMPT.clone(), max_new_tokens=budget) + + assert torch.equal( + output, plain_greedy_decode(target, PROMPT, max_new_tokens=budget) + ) + assert calls == {"begin": 1, "draft": 2, "verify": 2} + assert spec.step_trace[-1].accept == 4 + assert _cache_length(spec.last_target_cache) == PROMPT_LEN + budget + + +def _install_malformed_native_backend(spec: DFlashSpeculator) -> None: + target = spec.target + + class Backend: + def _native_model_forward_rich( + self, token_ids, metadata, logits_to_keep=0 + ): + if metadata is not None: + return (torch.zeros(1), torch.zeros(1)) + ids = torch.tensor([token_ids], dtype=torch.long) + outputs = target(ids, use_cache=True, output_hidden_states=True) + logits = ( + outputs.logits[:, -1:, :] if logits_to_keep else outputs.logits + ) + return logits, outputs.hidden_states, outputs.past_key_values + + spec.moe = Backend() + + +def test_malformed_backend_tuple_preserves_existing_error_and_no_trace() -> ( + None +): + spec = _tiny_spec() + _install_malformed_native_backend(spec) + + with pytest.raises( + RuntimeError, + match=r"_native_model_forward_rich must return \(logits, hidden_states, past_key_values\)", + ): + spec.generate(PROMPT.clone(), max_new_tokens=8) + + assert spec.step_trace == [] + assert spec.last_target_cache is None + assert spec.last_draft_cache is None + + +def test_drafted_tokens_are_not_exposed_before_successful_verification() -> ( + None +): + spec = _tiny_spec() + _install_malformed_native_backend(spec) + session = spec.begin_session(PROMPT.clone(), max_new_tokens=8) + verified_output = session.output_ids + + spec.draft_round(session) + + assert session.output_ids == verified_output + with pytest.raises(RuntimeError, match="must return"): + spec.verify_round(session) + assert session.output_ids == verified_output + assert spec.step_trace == [] + + +def test_hybrid_cache_rollback_runs_through_session_driver( + monkeypatch: pytest.MonkeyPatch, +) -> None: + qwen = pytest.importorskip( + "tests.python.dflash.test_qwen35_hybrid_rollback" + ) + target = qwen._tiny_qwen35_target(seed=9) + drafter = build_tiny_drafter( + target, + seed=10, + block_size=4, + target_layer_ids=(0, 1, 2, 3), + ) + config = read_dflash_config( + make_tiny_drafter_config( + target.config, + block_size=4, + target_layer_ids=(0, 1, 2, 3), + ) + ) + spec = DFlashSpeculator.from_models( + target, drafter, config=config, device="cpu" + ) + calls = _observe_session_driver(monkeypatch, spec) + + output = spec.generate(PROMPT.clone(), max_new_tokens=12) + + assert calls["begin"] == 1 + assert calls["draft"] == calls["verify"] == len(spec.step_trace) + assert any( + trace.accept + 1 < config.block_size for trace in spec.step_trace + ) + cached_length = _cache_length(spec.last_target_cache) + assert cached_length is not None + with torch.no_grad(): + expected = target( + output[:, :cached_length], use_cache=True + ).past_key_values + assert torch.allclose( + spec.last_target_cache.layers[0].conv_states, + expected.layers[0].conv_states, + ) + assert torch.allclose( + spec.last_target_cache.layers[0].recurrent_states, + expected.layers[0].recurrent_states, + ) diff --git a/tests/python/dflash/test_unified_protocols.py b/tests/python/dflash/test_unified_protocols.py new file mode 100644 index 00000000..082bf57a --- /dev/null +++ b/tests/python/dflash/test_unified_protocols.py @@ -0,0 +1,392 @@ +from __future__ import annotations + +from collections.abc import Callable +from dataclasses import fields +from typing import cast, final + +import pytest +import torch + +from moe_infinity.spec_decode._dflash_sample_ops import acceptance_sampled +from moe_infinity.spec_decode.protocols import ( + BackendCapabilities, + CacheAdapter, + CacheKind, + CacheSnapshot, + DenseCacheAdapter, + NativeStepTrace, + RequestSpec, + SamplingContext, + SessionRoundResult, + SessionTrace, +) + + +@final +class _FakeDenseCache: + def __init__(self, length: int) -> None: + self.length = length + self.crop_calls: list[int] = [] + + def get_seq_length(self) -> int: + return self.length + + def crop(self, length: int) -> None: + self.crop_calls.append(length) + self.length = length + + +def _step(*, accept: int = 2, emitted_len: int = 3) -> NativeStepTrace: + return NativeStepTrace( + prev_start=5, + accept=accept, + start=5 + accept + 1, + emitted_len=emitted_len, + target_cache_len=5 + accept + 1, + draft_cache_len=5, + ) + + +def _seed_ambient(seed: int) -> None: + manual_seed = cast(Callable[[int], torch.Generator], torch.manual_seed) + _ = manual_seed(seed) + + +def test_sampling_context_exposes_request_scoped_sampling_semantics() -> None: + generator = torch.Generator().manual_seed(17) + sampled = SamplingContext( + temperature=0.7, top_k=8, top_p=0.9, generator=generator + ) + greedy = SamplingContext(temperature=0.0) + + assert sampled.is_sampled + assert not sampled.is_greedy + assert sampled.generator is generator + assert greedy.is_greedy + assert not greedy.is_sampled + + +def test_sampling_context_rejects_invalid_values() -> None: + with pytest.raises(ValueError, match="temperature"): + _ = SamplingContext(temperature=-0.1) + with pytest.raises(ValueError, match="top_k"): + _ = SamplingContext(top_k=-1) + with pytest.raises(ValueError, match="top_p"): + _ = SamplingContext(top_p=0.0) + with pytest.raises(ValueError, match="top_p"): + _ = SamplingContext(top_p=1.1) + + +def test_request_spec_rejects_invalid_identity_prompt_and_budget() -> None: + with pytest.raises(ValueError, match="request_id"): + _ = RequestSpec("", (1,), 1) + with pytest.raises(ValueError, match="prompt_token_ids"): + _ = RequestSpec("req", (), 1) + with pytest.raises(ValueError, match="max_new_tokens"): + _ = RequestSpec("req", (1,), -1) + + +def test_request_spec_sampling_uses_a_default_factory() -> None: + sampling_field = next( + field for field in fields(RequestSpec) if field.name == "sampling" + ) + assert sampling_field.default_factory is SamplingContext + + +def test_request_and_round_contracts_have_stable_derived_counts() -> None: + request = RequestSpec( + request_id="req-7", + prompt_token_ids=(4, 5, 6), + max_new_tokens=12, + stop_token_ids=frozenset({2}), + sampling=SamplingContext(temperature=0.0), + ) + result = SessionRoundResult( + accepted_draft_count=2, + committed_token_ids=(7, 8, 9), + next_anchor=9, + target_cache_length=8, + emitted_length=3, + finished=False, + finish_reason=None, + ) + + assert request.prompt_length == 3 + assert request.is_sampled is False + assert result.cached_token_count == 3 + assert result.emitted_token_count == 3 + assert result.commit_block_token_ids == (7, 8, 9) + assert not hasattr(result, "emitted_token_ids") + + +def test_session_round_result_enforces_accept_plus_one_commit_invariant() -> ( + None +): + with pytest.raises(ValueError, match="accepted drafts plus one"): + _ = SessionRoundResult( + accepted_draft_count=2, + committed_token_ids=(7, 8), + next_anchor=8, + target_cache_length=8, + emitted_length=2, + finished=False, + finish_reason=None, + ) + + +def test_session_round_result_allows_a_finished_empty_noop() -> None: + result = SessionRoundResult( + accepted_draft_count=4, + committed_token_ids=(), + next_anchor=None, + target_cache_length=8, + emitted_length=0, + finished=True, + finish_reason="length", + ) + + assert result.cached_token_count == 0 + assert result.emitted_token_count == 0 + assert result.next_anchor is None + + +def test_session_round_result_rejects_an_unfinished_empty_commit() -> None: + with pytest.raises(ValueError, match="finished no-op"): + _ = SessionRoundResult( + accepted_draft_count=0, + committed_token_ids=(), + next_anchor=None, + target_cache_length=8, + emitted_length=0, + finished=False, + finish_reason=None, + ) + + +def test_session_round_result_exposes_the_approved_signature() -> None: + assert tuple(field.name for field in fields(SessionRoundResult)) == ( + "accepted_draft_count", + "committed_token_ids", + "next_anchor", + "target_cache_length", + "emitted_length", + "finished", + "finish_reason", + "fallback_reason", + ) + + +@pytest.mark.parametrize("cache_kind", ["dense_dynamic", "paged", "other"]) +def test_backend_capabilities_expose_the_approved_signature( + cache_kind: CacheKind, +) -> None: + capabilities = BackendCapabilities( + supports_batch=True, + supports_sampling=True, + supports_ragged_rows=False, + cache_kind=cache_kind, + supports_route_ahead=True, + supports_rich_forward=False, + ) + + assert tuple(field.name for field in fields(BackendCapabilities)) == ( + "supports_batch", + "supports_sampling", + "supports_ragged_rows", + "cache_kind", + "supports_route_ahead", + "supports_rich_forward", + "pairing_evidence", + "executor_evidence", + ) + assert capabilities.cache_kind == cache_kind + + +def test_backend_capabilities_reject_invalid_cache_kind() -> None: + invalid_kind = cast(CacheKind, cast(object, "dense")) + with pytest.raises(ValueError, match="cache_kind"): + _ = BackendCapabilities( + supports_batch=False, + supports_sampling=False, + supports_ragged_rows=False, + cache_kind=invalid_kind, + supports_route_ahead=False, + supports_rich_forward=False, + ) + + +@final +class _StructuralAdapter: + def snapshot(self) -> CacheSnapshot: + return CacheSnapshot(logical_length=0) + + def restore(self, snapshot: CacheSnapshot) -> None: + del snapshot + + def append(self, token_count: int) -> None: + del token_count + + def truncate(self, logical_length: int) -> None: + del logical_length + + def logical_length(self) -> int: + return 0 + + def release(self) -> None: + return None + + +def test_cache_adapter_is_a_runtime_structural_protocol() -> None: + assert getattr(CacheAdapter, "_is_protocol", False) + assert isinstance(_StructuralAdapter(), CacheAdapter) + + +def test_cache_snapshot_rejects_negative_logical_length() -> None: + with pytest.raises(ValueError, match="logical_length"): + _ = CacheSnapshot(logical_length=-1) + + +def test_dense_cache_adapter_implements_the_complete_cache_lifecycle() -> None: + cache = _FakeDenseCache(length=5) + adapter = DenseCacheAdapter(cache) + snapshot = adapter.snapshot() + + assert snapshot == CacheSnapshot(logical_length=5) + assert adapter.logical_length() == 5 + + cache.length = 8 + adapter.append(3) + assert adapter.logical_length() == 8 + + adapter.truncate(6) + assert adapter.logical_length() == 6 + assert cache.crop_calls == [6] + + snapshot = adapter.snapshot() + cache.length = 9 + adapter.append(3) + adapter.restore(snapshot) + assert adapter.logical_length() == 6 + assert cache.crop_calls == [6, 6] + + adapter.release() + assert cache.crop_calls == [6, 6, 0] + with pytest.raises(RuntimeError, match="released"): + _ = adapter.logical_length() + + +def test_dense_cache_adapter_rejects_invalid_length_transitions() -> None: + adapter = DenseCacheAdapter(_FakeDenseCache(length=3)) + + with pytest.raises(ValueError, match="token_count"): + adapter.append(-1) + with pytest.raises(ValueError, match="logical_length"): + adapter.truncate(4) + + +def test_session_trace_aggregates_native_step_trace_without_a_second_schema() -> ( + None +): + first = _step(accept=2, emitted_len=3) + second = NativeStepTrace( + prev_start=8, + accept=0, + start=9, + emitted_len=4, + target_cache_len=9, + draft_cache_len=8, + ) + trace = SessionTrace( + request_id="req-7", + backend="native", + cache_kind="dense_dynamic", + sampled=False, + route_ahead_status="disabled", + ) + trace.append(first) + trace.append(second) + trace.rollback = 1 + trace.replay = 1 + trace.finish_reason = "length" + + assert trace.round_count == 2 + assert trace.accepted == 2 + assert trace.committed == 4 + assert trace.emitted == 4 + assert first.committed_count == 3 + assert trace.as_dict() == { + "request_id": "req-7", + "backend": "native", + "cache_kind": "dense_dynamic", + "sampled": False, + "round_count": 2, + "accepted": 2, + "committed": 4, + "emitted": 4, + "rollback": 1, + "replay": 1, + "finish_reason": "length", + "route_ahead_status": "disabled", + "pairing_evidence": { + "valid": False, + "config_valid": False, + "dimensions_valid": False, + "vocab_valid": False, + "mask_valid": False, + "layers_valid": False, + "block_valid": False, + "module_valid": None, + "validated_checkpoint_scope": (), + "failure_reason": None, + }, + "executor_evidence": { + "wiring_reachable": False, + "prefetcher_present": False, + "attempted_layers": (), + "fired_layers": (), + "actual_expert_union": (), + "prefetched_bytes": 0, + "coverage": None, + "wasted_prefetch_bytes": None, + "cache_hit_rate": None, + "fallback_reason": None, + }, + } + + +def test_acceptance_sampled_uses_only_the_explicit_generator() -> None: + draft_probs = torch.tensor([[0.5, 0.5]]) + target_probs = torch.tensor([[0.2, 0.8], [0.7, 0.3]]) + drafts = torch.tensor([0]) + + _seed_ambient(999) + ambient_before = torch.random.get_rng_state().clone() + first = acceptance_sampled( + draft_probs, + target_probs, + drafts, + generator=torch.Generator().manual_seed(123), + ) + ambient_after = torch.random.get_rng_state() + second = acceptance_sampled( + draft_probs, + target_probs, + drafts, + generator=torch.Generator().manual_seed(123), + ) + + assert first == second + assert torch.equal(ambient_before, ambient_after) + + +def test_acceptance_sampled_keeps_ambient_rng_compatibility() -> None: + draft_probs = torch.tensor([[0.5, 0.5]]) + target_probs = torch.tensor([[0.2, 0.8], [0.7, 0.3]]) + drafts = torch.tensor([0]) + + _seed_ambient(321) + first = acceptance_sampled(draft_probs, target_probs, drafts) + _seed_ambient(321) + second = acceptance_sampled(draft_probs, target_probs, drafts) + + assert first == second diff --git a/tests/python/integration/test_deepseek_mla_adapter.py b/tests/python/integration/test_deepseek_mla_adapter.py new file mode 100644 index 00000000..a0130094 --- /dev/null +++ b/tests/python/integration/test_deepseek_mla_adapter.py @@ -0,0 +1,372 @@ +from __future__ import annotations + +import copy + +import pytest +import torch + +transformers = pytest.importorskip("transformers") + +from moe_infinity.models.deepseek_mla_attention import ( # noqa: E402 + adapt_deepseek_attention, + adapt_deepseek_model, + clear_deepseek_mla_context, + is_deepseek_mla_eligible, + set_deepseek_mla_context, +) +from moe_infinity.runtime.attention_types import AttentionMetadata # noqa: E402 +from moe_infinity.serving.mla_cache import MLAPagedKVCache # noqa: E402 +from moe_infinity.spec_decode.protocols import RichForwardResult # noqa: E402 + + +def _metadata( + *, + seq_len: int, + slots: list[int], + prefill: bool, + block_table: list[int] | None = None, + seq_id: int | None = None, +) -> AttentionMetadata: + metadata = AttentionMetadata( + block_tables=torch.tensor( + [block_table if block_table is not None else [0, 1]], + dtype=torch.int32, + ), + seq_lens=torch.tensor([seq_len], dtype=torch.int32), + max_seq_len=seq_len, + num_prefill_tokens=len(slots) if prefill else 0, + num_decode_tokens=0 if prefill else len(slots), + slot_mapping=torch.tensor(slots, dtype=torch.int64), + is_prefill=prefill, + ) + metadata.seq_id = seq_id + return metadata + + +def _case( + version: str, + *, + q_lora_rank: int | None = None, + rope_interleave: bool = True, +): + if version == "v2": + modeling = pytest.importorskip( + "transformers.models.deepseek_v2.modeling_deepseek_v2" + ) + config_cls = transformers.DeepseekV2Config + attention_cls = modeling.DeepseekV2Attention + rotary_cls = modeling.DeepseekV2RotaryEmbedding + config = config_cls( + hidden_size=16, + intermediate_size=24, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=q_lora_rank, + kv_lora_rank=4, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=4, + head_dim=8, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + moe_intermediate_size=8, + ) + else: + modeling = pytest.importorskip( + "transformers.models.deepseek_v3.modeling_deepseek_v3" + ) + config_cls = transformers.DeepseekV3Config + attention_cls = modeling.DeepseekV3Attention + rotary_cls = modeling.DeepseekV3RotaryEmbedding + config = config_cls( + hidden_size=16, + intermediate_size=24, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=q_lora_rank, + kv_lora_rank=4, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=4, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + moe_intermediate_size=8, + n_group=1, + topk_group=1, + rope_interleave=rope_interleave, + ) + return modeling, config, attention_cls, rotary_cls + + +def _positions( + version: str, + rotary: torch.nn.Module, + hidden: torch.Tensor, + positions: torch.Tensor, +): + if version == "v2": + return rotary(hidden, positions) + return rotary(hidden, positions) + + +def _call( + version: str, + module: torch.nn.Module, + hidden: torch.Tensor, + positions: object, + cache: object = None, + attention_mask: torch.Tensor | None = None, +) -> torch.Tensor: + kwargs = dict( + hidden_states=hidden, + attention_mask=attention_mask, + past_key_values=cache, + position_embeddings=positions, + ) + return module(**kwargs)[0] + + +def _validation_case(): + _, config, attention_cls, rotary_cls = _case("v3") + module = attention_cls(config, layer_idx=0).eval() + cache = MLAPagedKVCache(4, 2, 1, 4, 4, torch.float32, torch.device("cpu")) + cache.allocate_sequence(17, 3) + adapt_deepseek_attention(module, cache, enabled=True) + hidden = torch.randn(1, 3, config.hidden_size) + positions = rotary_cls(config)(hidden, torch.arange(3).unsqueeze(0)) + return module, cache, hidden, positions + + +def test_mla_attention_rejects_total_len_shorter_than_query() -> None: + module, _, hidden, positions = _validation_case() + set_deepseek_mla_context( + module, _metadata(seq_len=2, slots=[0, 1, 2], prefill=True) + ) + try: + with pytest.raises(ValueError, match="total_len .* query_len"): + _call("v3", module, hidden, positions) + finally: + clear_deepseek_mla_context(module) + + +def test_mla_attention_rejects_short_attention_mask() -> None: + module, _, hidden, positions = _validation_case() + set_deepseek_mla_context( + module, _metadata(seq_len=3, slots=[0, 1, 2], prefill=True) + ) + try: + with pytest.raises(ValueError, match="attention_mask last dimension"): + _call( + "v3", + module, + hidden, + positions, + attention_mask=torch.zeros(1, 1, 3, 2), + ) + finally: + clear_deepseek_mla_context(module) + + +def test_mla_attention_requires_seq_id_for_engine_owned_cache_access() -> None: + module, _, hidden, positions = _validation_case() + set_deepseek_mla_context( + module, _metadata(seq_len=3, slots=[0, 1, 2], prefill=True) + ) + try: + with pytest.raises(ValueError, match="requires metadata.seq_id"): + _call("v3", module, hidden, positions) + finally: + clear_deepseek_mla_context(module) + + +@pytest.mark.parametrize( + ("block_table", "slots", "message"), + [ + ([1, 0], [0, 1, 2], "block_tables do not match allocated sequence"), + ( + [0, 1], + [0, 1, 4], + "slot_mapping references pages not owned by sequence", + ), + ], +) +def test_mla_attention_rejects_incoherent_owned_pages( + block_table: list[int], slots: list[int], message: str +) -> None: + module, cache, hidden, positions = _validation_case() + set_deepseek_mla_context( + module, + _metadata( + seq_len=3, + slots=slots, + prefill=True, + block_table=block_table, + seq_id=17, + ), + ) + try: + with pytest.raises(ValueError, match=message): + _call("v3", module, hidden, positions) + finally: + clear_deepseek_mla_context(module) + + +@pytest.mark.parametrize( + ("version", "q_lora_rank", "rope_interleave"), + [ + pytest.param("v2", None, True, id="v2-direct-q"), + pytest.param("v2", 16, True, id="v2-lora-q"), + pytest.param("v3", None, True, id="v3-direct-q-interleaved"), + pytest.param("v3", 16, True, id="v3-lora-q-interleaved"), + pytest.param("v3", None, False, id="v3-direct-q-split-rope"), + pytest.param("v3", 16, False, id="v3-lora-q-split-rope"), + ], +) +def test_real_upstream_attention_is_adapted_in_place_and_matches_dense( + version: str, + q_lora_rank: int | None, + rope_interleave: bool, +) -> None: + modeling, config, attention_cls, rotary_cls = _case( + version, + q_lora_rank=q_lora_rank, + rope_interleave=rope_interleave, + ) + torch.manual_seed(4) + dense = attention_cls(config, layer_idx=0).eval() + paged = copy.deepcopy(dense) + parameter_ids = { + name: id(value) for name, value in paged.named_parameters() + } + cache = MLAPagedKVCache(4, 2, 1, 4, 4, torch.float32, torch.device("cpu")) + cache.allocate_sequence(17, 3) + + adapted = adapt_deepseek_attention(paged, cache, enabled=True) + + assert adapted is paged + assert adapted.layer_idx == 0 + assert parameter_ids == { + name: id(value) for name, value in adapted.named_parameters() + } + assert ( + adapted.__class__.__name__ + == f"Deepseek{version.upper()}MLAPagedAttention" + ) + + rotary = rotary_cls(config) + prompt = torch.randn(1, 3, config.hidden_size) + prompt_pos = torch.arange(3).unsqueeze(0) + prompt_embeddings = _positions(version, rotary, prompt, prompt_pos) + prompt_mask = torch.zeros(1, 1, 3, 3) + prompt_mask.masked_fill_( + torch.triu(torch.ones(3, 3, dtype=torch.bool), diagonal=1), + torch.finfo(prompt.dtype).min, + ) + dense_prompt = _call( + version, dense, prompt, prompt_embeddings, attention_mask=prompt_mask + ) + + set_deepseek_mla_context( + adapted, + _metadata( + seq_len=3, + slots=[0, 1, 2], + prefill=True, + seq_id=17, + ), + ) + try: + paged_prompt = _call( + version, + adapted, + prompt, + prompt_embeddings, + attention_mask=prompt_mask, + ) + finally: + clear_deepseek_mla_context(adapted) + + assert torch.allclose(paged_prompt, dense_prompt, atol=2e-5, rtol=2e-5) + assert torch.count_nonzero(cache.get_mla_cache_tensors()[0]) > 0 + + from transformers import DynamicCache + + dense_cache = DynamicCache() + _ = _call( + version, + dense, + prompt, + prompt_embeddings, + dense_cache, + prompt_mask, + ) + token = torch.randn(1, 1, config.hidden_size) + decode_pos = torch.tensor([[3]]) + decode_embeddings = _positions(version, rotary, token, decode_pos) + dense_decode = _call(version, dense, token, decode_embeddings, dense_cache) + + cache.append_tokens(17, 1) + set_deepseek_mla_context( + adapted, + _metadata(seq_len=4, slots=[3], prefill=False, seq_id=17), + ) + try: + paged_decode = _call(version, adapted, token, decode_embeddings) + finally: + clear_deepseek_mla_context(adapted) + + assert torch.allclose(paged_decode, dense_decode, atol=2e-5, rtol=2e-5) + assert cache is adapted._mla_cache + + +def test_selection_is_default_off_and_rejects_hybrid_models() -> None: + from moe_infinity.utils.config import ArcherConfig + + assert ArcherConfig().enable_deepseek_mla_paging is False + _, config, attention_cls, _ = _case("v3") + module = attention_cls(config, layer_idx=0) + cache = MLAPagedKVCache(2, 2, 1, 4, 4, torch.float32, torch.device("cpu")) + + assert adapt_deepseek_attention(module, cache) is module + assert module.__class__ is attention_cls + assert not is_deepseek_mla_eligible(config, enabled=False) + config.sliding_window = 32 + assert not is_deepseek_mla_eligible(config, enabled=True) + + +def test_native_rich_forward_returns_engine_cache_handle_without_dynamic_cache() -> ( + None +): + from moe_infinity.entrypoints.big_modeling import MoE + + _, config, _, _ = _case("v3") + config.vocab_size = 32 + config.first_k_dense_replace = 1 + model_cls = transformers.DeepseekV3ForCausalLM + model = model_cls(config).eval() + cache = MLAPagedKVCache(4, 2, 1, 4, 4, torch.float32, torch.device("cpu")) + modules = adapt_deepseek_model(model, cache, enabled=True) + assert len(modules) == 1 + + shell = MoE.__new__(MoE) + shell.model = model + shell._cached_past_key_values = None + shell._native_attention_backend = None + shell._native_mla_cache = cache + shell._resolve_native_input_device = lambda: torch.device("cpu") + cache.allocate_sequence(23, 3) + metadata = _metadata(seq_len=3, slots=[0, 1, 2], prefill=True, seq_id=23) + draft_cache = object() + + result = shell._native_model_forward_rich([1, 2, 3], metadata) + + assert isinstance(result, RichForwardResult) + assert result.logits.shape == (1, 3, config.vocab_size) + assert len(result.hidden_states) == 2 + assert result.cache_handle is cache + assert result.cache_handle is not draft_cache + assert shell._cached_past_key_values is None diff --git a/tests/python/integration/test_dflash_unified_validation_e2e.py b/tests/python/integration/test_dflash_unified_validation_e2e.py new file mode 100644 index 00000000..7ac86ff1 --- /dev/null +++ b/tests/python/integration/test_dflash_unified_validation_e2e.py @@ -0,0 +1,70 @@ +from __future__ import annotations + +import json +import os +import subprocess +import sys +from pathlib import Path +from typing import Any + +ROOT = Path(__file__).resolve().parents[3] +VALIDATOR = ROOT / "benchmarks/dflash/validate_unified_execution.py" + + +def _run_validator( + working_directory: Path, *arguments: str +) -> subprocess.CompletedProcess[str]: + return subprocess.run( + [sys.executable, str(VALIDATOR), "--fixture", "tiny", *arguments], + cwd=working_directory, + check=False, + capture_output=True, + text=True, + env={**os.environ, "MOE_DFLASH_GPU": "0"}, + ) + + +def _report(completed: subprocess.CompletedProcess[str]) -> dict[str, Any]: + parsed = json.loads(completed.stdout) + assert isinstance(parsed, dict) + return parsed + + +def test_validator_passes_cpu_rollout_gates_from_non_repo_cwd( + tmp_path: Path, +) -> None: + completed = _run_validator( + tmp_path, + "--require-cache-invariants", + "--require-order-invariance", + "--json", + ) + + assert completed.returncode == 0, completed.stderr + report = _report(completed) + assert report["status"] == "PASS" + assert ( + report["compatibility"]["execution_mode"] == "tiny_cpu_protocol_fixture" + ) + assert report["cache_invariants"] is True + assert report["order_invariance"] is True + assert report["ownership_isolation"] is True + assert report["paged_ownership_released"] is True + assert report["pairing_executor_separate"] is True + assert report["checkpoint_downloads"] is False + assert report["gpu_harness_executed"] is False + + +def test_validator_require_gpu_fails_closed_when_fixture_is_disabled( + tmp_path: Path, +) -> None: + completed = _run_validator(tmp_path, "--require-gpu", "--json") + + assert completed.returncode == 1, completed.stderr + report = _report(completed) + assert report["status"] == "FAIL" + assert report["gpu_readiness_required"] is True + assert report["gpu_fixture_enabled"] is False + assert report["gpu_readiness_pass"] is False + assert report["checkpoint_downloads"] is False + assert report["gpu_harness_executed"] is False diff --git a/tests/python/serving/test_adaptive_scheduler.py b/tests/python/serving/test_adaptive_scheduler.py index f4dc380c..bd896294 100644 --- a/tests/python/serving/test_adaptive_scheduler.py +++ b/tests/python/serving/test_adaptive_scheduler.py @@ -2,16 +2,16 @@ import sys import time from pathlib import Path +from types import ModuleType from typing import Optional, Protocol, cast ROOT = str(Path(__file__).resolve().parents[3]) if ROOT not in sys.path: sys.path.insert(0, ROOT) -_ = sys.modules.pop("moe_infinity", None) -_ = sys.modules.pop("moe_infinity.serving", None) MEMORY_MANAGER_PATH = ( Path(ROOT) / "moe_infinity" / "serving" / "memory_manager.py" ) +_MISSING_MODULE = object() class AdaptiveKVSchedulerProtocol(Protocol): @@ -46,8 +46,15 @@ def _load_adaptive_scheduler() -> type[AdaptiveKVSchedulerProtocol]: if spec is None or spec.loader is None: raise RuntimeError(f"failed to load module from {MEMORY_MANAGER_PATH}") module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) + previous_module = sys.modules.get(module_name, _MISSING_MODULE) + try: + sys.modules[module_name] = module + spec.loader.exec_module(module) + finally: + if previous_module is _MISSING_MODULE: + _ = sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = cast(ModuleType, previous_module) return cast( type[AdaptiveKVSchedulerProtocol], getattr(module, "AdaptiveKVScheduler"), diff --git a/tests/python/serving/test_api_routes.py b/tests/python/serving/test_api_routes.py index 6b63e7b7..8fc45423 100644 --- a/tests/python/serving/test_api_routes.py +++ b/tests/python/serving/test_api_routes.py @@ -1,6 +1,7 @@ # pyright: reportAny=false, reportCallIssue=false, reportExplicitAny=false, reportMissingParameterType=false, reportMissingTypeArgument=false, reportUnknownArgumentType=false, reportUnknownMemberType=false, reportUnknownParameterType=false, reportUnknownVariableType=false from __future__ import annotations +import sys from dataclasses import dataclass, field from types import SimpleNamespace from typing import Any, cast @@ -186,11 +187,80 @@ def _capture_engine(**kwargs: Any) -> MagicMock: tok=None, max_seq_length=128, speculative_draft=speculator, + enable_deepseek_mla_paging=True, + max_resident_paged_speculative_sessions=3, + min_free_mla_blocks_after_admission=2, ) assert captured["model"] is model assert captured["engine"] is offload_engine assert captured["speculative_draft"] is speculator + assert captured["config"]["enable_deepseek_mla_paging"] is True + assert captured["config"]["max_resident_paged_speculative_sessions"] == 3 + assert captured["config"]["min_free_mla_blocks_after_admission"] == 2 + + +def test_dflash_paged_cli_defaults_remain_off_and_bounded( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + sys, + "argv", + [ + "api_server_v2.py", + "--model", + "demo/model", + "--offload-dir", + "/tmp/offload", + ], + ) + + args = srv.parse_args() + + assert args.enable_deepseek_mla_paging is False + assert args.max_resident_paged_speculative_sessions == 1 + assert args.min_free_mla_blocks_after_admission == 1 + + +def test_dflash_paged_cli_values_forward_to_engine_config( + monkeypatch: pytest.MonkeyPatch, +) -> None: + monkeypatch.setattr( + sys, + "argv", + [ + "api_server_v2.py", + "--model", + "demo/model", + "--offload-dir", + "/tmp/offload", + "--enable-deepseek-mla-paging", + "--max-resident-paged-speculative-sessions", + "4", + "--min-free-mla-blocks-after-admission", + "3", + ], + ) + args = srv.parse_args() + model = SimpleNamespace( + config=SimpleNamespace( + num_hidden_layers=2, + num_attention_heads=4, + num_key_value_heads=2, + hidden_size=32, + head_dim=8, + max_position_embeddings=128, + eos_token_id=2, + dtype="float32", + ), + dtype="float32", + ) + + config = srv._build_engine_config(args=args, model=model) + + assert config["enable_deepseek_mla_paging"] is True + assert config["max_resident_paged_speculative_sessions"] == 4 + assert config["min_free_mla_blocks_after_admission"] == 3 def test_list_models_engine_not_ready(client: TestClient) -> None: diff --git a/tests/python/serving/test_dflash_stage4a.py b/tests/python/serving/test_dflash_stage4a.py new file mode 100644 index 00000000..f770d1dc --- /dev/null +++ b/tests/python/serving/test_dflash_stage4a.py @@ -0,0 +1,580 @@ +from __future__ import annotations + +import logging +import types +from dataclasses import dataclass, field +from typing import Callable + +import pytest +import torch + +from moe_infinity.serving.engine import ContinuousBatchingEngine +from moe_infinity.serving.sequence import SamplingParams, SequenceStatus +from moe_infinity.serving.spec_session_driver import ( + EXECUTION_CONTEXT_TEMPORARY_DYNAMIC, + SpecSessionDriver, +) + + +class _Cache: + def __init__(self, *, fail_crop: bool = False) -> None: + self.released = False + self.fail_crop = fail_crop + self.crop_calls = 0 + self.crop_error = RuntimeError("cache cleanup boom") + + def crop(self, length: int) -> None: + self.crop_calls += 1 + if self.fail_crop: + raise self.crop_error + if length == 0: + self.released = True + + +@dataclass +class _Session: + emitted: list[int] + max_new_tokens: int + target_kv: _Cache = field(default_factory=_Cache) + draft_kv: _Cache = field(default_factory=_Cache) + finished: bool = False + pending: bool = False + + @property + def output_ids(self) -> list[int]: + return self.emitted[: self.max_new_tokens] + + def clear_pending(self) -> None: + self.pending = False + + +class _Speculator: + def __init__(self) -> None: + self.moe = types.SimpleNamespace(_cached_past_key_values=None) + self.begin_calls: list[dict[str, object]] = [] + self.sessions: list[_Session] = [] + self.events: list[tuple[str, int]] = [] + self.verify_hook: Callable[[], None] | None = None + self.replace_target_cache = False + self.draft_error: BaseException | None = None + self.verify_error: BaseException | None = None + + def begin_session( + self, input_ids: torch.Tensor, **kwargs: object + ) -> _Session: + self.begin_calls.append(dict(kwargs)) + anchor = int(input_ids[0, -1]) + 1 + session = _Session( + emitted=[anchor], max_new_tokens=int(kwargs["max_new_tokens"]) + ) + session.finished = session.max_new_tokens <= 1 + self.sessions.append(session) + self.events.append(("begin", anchor)) + return session + + def draft_round(self, session: _Session) -> object: + assert self.moe._cached_past_key_values is session.target_kv + if self.draft_error is not None: + raise self.draft_error + session.pending = True + self.events.append(("draft", session.emitted[-1])) + return types.SimpleNamespace(tokens=4, expert_bytes=64) + + def verify_round(self, session: _Session) -> object: + assert self.moe._cached_past_key_values is session.target_kv + assert session.pending + if self.verify_error is not None: + raise self.verify_error + self.events.append(("verify", session.emitted[-1])) + if self.verify_hook is not None: + self.verify_hook() + if self.replace_target_cache: + session.target_kv = _Cache() + session.pending = False + remaining = session.max_new_tokens - len(session.emitted) + for _ in range(min(2, remaining)): + session.emitted.append(session.emitted[-1] + 1) + session.finished = len(session.emitted) >= session.max_new_tokens + return types.SimpleNamespace( + accepted_token_ids=session.emitted[-min(2, remaining) :], + committed_count=min(2, remaining), + finished=session.finished, + ) + + +class _Model: + config = types.SimpleNamespace(vocab_size=128, eos_token_id=99) + + def eval(self) -> None: + pass + + def forward( + self, input_ids: torch.Tensor, **kwargs: object + ) -> types.SimpleNamespace: + del kwargs + logits = torch.full((*input_ids.shape, 128), -1e9) + for row in range(input_ids.shape[0]): + for col in range(input_ids.shape[1]): + logits[row, col, int(input_ids[row, col]) + 1] = 0 + return types.SimpleNamespace(logits=logits) + + +class _Offload: + def __init__(self) -> None: + self.request_id = 0 + self.expert_tracer = types.SimpleNamespace(create_entry=lambda: 0) + self.expert_layer_modules = [types.SimpleNamespace(seq_id_list=[])] + + def _generate_request_id(self) -> int: + value = self.request_id + self.request_id += 1 + return value + + +def _config(**overrides: object) -> dict[str, object]: + config: dict[str, object] = { + "device_memory_ratio": 0.75, + "kv_cache_ratio": 0.25, + "max_batch_size": 8, + "max_tokens_per_step": 16, + "block_size": 4, + "num_layers": 1, + "num_kv_heads": 2, + "head_dim": 8, + "dtype": "float32", + "eos_token_id": 99, + "num_kv_blocks": 32, + "verify_token_budget": 8, + "verify_expert_byte_budget": 128, + "verify_token_deficit_cap": 32, + "verify_expert_byte_deficit_cap": 512, + } + config.update(overrides) + return config + + +def _engine(speculator: object, **config: object) -> ContinuousBatchingEngine: + return ContinuousBatchingEngine( + model=_Model(), + engine=_Offload(), + config=_config(**config), + speculative_draft=speculator, + ) + + +def test_driver_record_owns_temporary_context_and_logical_commit_state() -> ( + None +): + speculator = _Speculator() + driver = SpecSessionDriver(speculator) + + record = driver.begin( + request_id="r0", + seq_id=7, + prompt_token_ids=[3, 4], + max_new_tokens=3, + temperature=0.7, + top_k=5, + top_p=0.8, + stop_token_ids=[99], + callbacks=(), + ) + + assert record.request_id == "r0" and record.seq_id == 7 + assert record.spec_session is speculator.sessions[0] + assert record.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert record.decode_state.invariant_holds() + assert driver.commit(record) == (5,) + assert record.decode_state.cached_len == 3 + assert record.decode_state.invariant_holds() + + +def test_driver_tracks_replaced_private_target_cache_for_cleanup() -> None: + speculator = _Speculator() + speculator.replace_target_cache = True + driver = SpecSessionDriver(speculator) + record = driver.begin( + request_id="replace", + seq_id=8, + prompt_token_ids=[3], + max_new_tokens=3, + temperature=0.0, + top_k=0, + top_p=1.0, + stop_token_ids=[99], + callbacks=(), + ) + original = record.spec_session.target_kv + _ = driver.commit(record) + _ = driver.draft(record) + _ = driver.verify(record) + replacement = record.spec_session.target_kv + + assert record.execution_context.target_cache is replacement + assert original.released + assert not replacement.released + assert not record.spec_session.draft_kv.released + assert len(record.execution_context.owned_caches) == 2 + driver.release(record) + assert replacement.released + + +def test_engine_creates_and_persists_one_session_per_sequence() -> None: + speculator = _Speculator() + engine = _engine(speculator) + engine.add_request( + "multi", + [10], + SamplingParams(temperature=0.7, top_p=0.9, max_tokens=5), + n=2, + ) + + first = engine.step() + + assert [output.token_id for output in first] == [11, 11] + assert len(speculator.sessions) == 2 + assert len(engine.speculative_sessions) == 2 + assert all( + engine._sequences[seq_id].status is SequenceStatus.DRAFT + for seq_id in engine._request_to_seq_ids["multi"] + ) + + second = engine.step() + assert [output.token_id for output in second] == [12, 13, 12, 13] + assert len(engine.speculative_sessions) == 2 + + +def test_sampled_parameters_stops_and_callbacks_are_preserved() -> None: + speculator = _Speculator() + engine = _engine(speculator) + streamed: list[int] = [] + engine.add_request( + "sampled", + [20], + SamplingParams(temperature=0.75, top_k=7, top_p=0.85, max_tokens=4), + on_token=lambda output: streamed.append(output.token_id), + ) + + result = engine.run_until_done() + + assert result == {"sampled": [21, 22, 23, 24]} + assert streamed == [21, 22, 23, 24] + assert speculator.begin_calls[0]["temperature"] == 0.75 + assert speculator.begin_calls[0]["top_k"] == 7 + assert speculator.begin_calls[0]["top_p"] == 0.85 + assert speculator.begin_calls[0]["stop_token_ids"] == [99] + + +def test_cancellation_during_verify_suppresses_output_and_releases_caches() -> ( + None +): + speculator = _Speculator() + engine = _engine(speculator) + streamed: list[int] = [] + engine.add_request( + "cancel", + [30], + SamplingParams(temperature=0.6, max_tokens=5), + on_token=lambda output: streamed.append(output.token_id), + ) + assert [output.token_id for output in engine.step()] == [31] + session = speculator.sessions[0] + speculator.verify_hook = lambda: engine.abort_request("cancel") + + assert engine.step() == [] + assert streamed == [31] + assert session.target_kv.released and session.draft_kv.released + assert engine.speculative_sessions == {} + assert engine.has_pending_requests() is False + + +def test_draft_failure_cleans_request_without_streaming_unverified_tokens() -> ( + None +): + speculator = _Speculator() + engine = _engine(speculator) + streamed: list[int] = [] + engine.add_request( + "draft-fail", + [30], + SamplingParams(temperature=0.5, max_tokens=5), + on_token=lambda output: streamed.append(output.token_id), + ) + assert [output.token_id for output in engine.step()] == [31] + session = speculator.sessions[0] + original = RuntimeError("draft secret at /srv/private/model.bin") + speculator.draft_error = original + + with pytest.raises(RuntimeError, match="draft secret") as caught: + engine.step() + + assert caught.value is original + assert streamed == [31] + assert session.target_kv.released and session.draft_kv.released + assert engine.speculative_sessions == {} + assert engine.scheduler._verify_demands == {} + assert engine.has_pending_requests() is False + assert engine.get_request_failure("draft-fail") == { + "phase": "draft", + "failure_type": "RuntimeError", + "code": "speculative_draft_failed", + } + + +def test_verify_failure_cleans_pending_draft_without_streaming_it() -> None: + speculator = _Speculator() + engine = _engine(speculator) + streamed: list[int] = [] + engine.add_request( + "verify-fail", + [40], + SamplingParams(temperature=0.5, max_tokens=5), + on_token=lambda output: streamed.append(output.token_id), + ) + assert [output.token_id for output in engine.step()] == [41] + session = speculator.sessions[0] + original = ValueError("verify boom") + speculator.verify_error = original + + with pytest.raises(ValueError, match="verify boom") as caught: + engine.step() + + assert caught.value is original + assert streamed == [41] + assert session.pending is False + assert session.target_kv.released and session.draft_kv.released + assert engine.speculative_sessions == {} + assert engine.scheduler._verify_demands == {} + assert engine.has_pending_requests() is False + assert engine.get_request_failure("verify-fail")["phase"] == "verify" + + +def test_backend_failure_survives_private_cache_cleanup_failure() -> None: + speculator = _Speculator() + engine = _engine(speculator) + engine.add_request( + "cleanup-fail", + [50], + SamplingParams(temperature=0.5, max_tokens=5), + ) + _ = engine.step() + session = speculator.sessions[0] + session.target_kv.fail_crop = True + original = RuntimeError("draft primary") + speculator.draft_error = original + + with pytest.raises(RuntimeError, match="draft primary") as caught: + engine.step() + + assert caught.value is original + assert session.target_kv.crop_calls == 1 + assert session.draft_kv.released + assert engine.speculative_sessions == {} + assert engine.has_pending_requests() is False + assert caught.value.session_cleanup_errors == ( + session.target_kv.crop_error, + ) + notes = getattr(caught.value, "__notes__", None) + if notes is not None: + assert any("cache cleanup boom" in note for note in notes) + assert engine.get_request_failure("cleanup-fail") == { + "phase": "draft", + "failure_type": "RuntimeError", + "code": "speculative_draft_failed", + } + + +def test_cleanup_metadata_does_not_depend_on_add_note() -> None: + class LegacyRuntimeError(RuntimeError): + def __getattribute__(self, name: str) -> object: + if name == "add_note": + return None + return super().__getattribute__(name) + + speculator = _Speculator() + engine = _engine(speculator) + engine.add_request( + "legacy-cleanup-fail", + [50], + SamplingParams(temperature=0.5, max_tokens=5), + ) + _ = engine.step() + session = speculator.sessions[0] + session.target_kv.fail_crop = True + original = LegacyRuntimeError("draft primary") + speculator.draft_error = original + + with pytest.raises(LegacyRuntimeError, match="draft primary") as caught: + engine.step() + + assert caught.value is original + assert caught.value.session_cleanup_errors == ( + session.target_kv.crop_error, + ) + + +def test_cleanup_reporting_failures_are_logged_without_masking_primary( + caplog: pytest.LogCaptureFixture, +) -> None: + class HostileRuntimeError(RuntimeError): + def __setattr__(self, name: str, value: object) -> None: + if name == "session_cleanup_errors": + raise TypeError("metadata blocked") + super().__setattr__(name, value) + + def add_note(self, note: str) -> None: + del note + raise RuntimeError("note blocked") + + speculator = _Speculator() + engine = _engine(speculator) + engine.add_request( + "hostile-cleanup-fail", + [50], + SamplingParams(temperature=0.5, max_tokens=5), + ) + _ = engine.step() + session = speculator.sessions[0] + session.target_kv.fail_crop = True + original = HostileRuntimeError("draft primary") + speculator.draft_error = original + + with ( + caplog.at_level(logging.DEBUG, logger="moe_infinity.serving.engine"), + pytest.raises(HostileRuntimeError, match="draft primary") as caught, + ): + engine.step() + + assert caught.value is original + assert session.target_kv.crop_calls == 1 + assert session.draft_kv.released + assert engine.speculative_sessions == {} + assert engine.has_pending_requests() is False + assert any( + "cleanup metadata attachment failed for request " + "hostile-cleanup-fail during draft; " + "primary=HostileRuntimeError reporting=TypeError" in message + for message in caplog.messages + ) + assert any( + "cleanup note attachment failed for request " + "hostile-cleanup-fail during draft; " + "primary=HostileRuntimeError reporting=RuntimeError" in message + for message in caplog.messages + ) + + +def test_cleanup_note_lookup_failure_is_logged_without_masking_primary( + caplog: pytest.LogCaptureFixture, +) -> None: + class LookupHostileRuntimeError(RuntimeError): + def __getattribute__(self, name: str) -> object: + if name == "add_note": + raise LookupError("note lookup blocked") + return super().__getattribute__(name) + + speculator = _Speculator() + engine = _engine(speculator) + engine.add_request( + "lookup-cleanup-fail", + [50], + SamplingParams(temperature=0.5, max_tokens=5), + ) + _ = engine.step() + session = speculator.sessions[0] + session.target_kv.fail_crop = True + original = LookupHostileRuntimeError("draft primary") + record = next(iter(engine.speculative_sessions.values())) + + with ( + caplog.at_level(logging.DEBUG, logger="moe_infinity.serving.engine"), + pytest.raises( + LookupHostileRuntimeError, match="draft primary" + ) as caught, + ): + engine._fail_speculative_request(record, "draft", original) + raise original + + assert caught.value is original + assert caught.value.session_cleanup_errors == ( + session.target_kv.crop_error, + ) + assert session.draft_kv.released + assert engine.speculative_sessions == {} + assert engine.has_pending_requests() is False + assert any( + "cleanup note-lookup attachment failed for request " + "lookup-cleanup-fail during draft; " + "primary=LookupHostileRuntimeError reporting=LookupError" in message + for message in caplog.messages + ) + + +def test_mixed_scheduled_work_splits_eligible_from_normal_fallback() -> None: + speculator = _Speculator() + engine = _engine(speculator) + engine.add_request( + "eligible", + [40], + SamplingParams(temperature=0.7, max_tokens=3), + ) + engine.add_request( + "fallback", + [50], + SamplingParams(temperature=0.0, repetition_penalty=1.1, max_tokens=1), + ) + + outputs = engine.step() + + assert [(row.request_id, row.token_id) for row in outputs] == [ + ("eligible", 41), + ("fallback", 51), + ] + assert len(speculator.sessions) == 1 + assert ( + engine._sequences[engine._request_to_seq_ids["eligible"][0]].status + is SequenceStatus.DRAFT + ) + + +def test_grammar_metadata_keeps_request_on_normal_serving_fallback() -> None: + speculator = _Speculator() + engine = _engine(speculator) + params = SamplingParams(temperature=0.0, max_tokens=1) + params.grammar = "root ::= 'x'" # type: ignore[attr-defined] + engine.add_request("grammar", [70], params) + + outputs = engine.step() + + assert [output.token_id for output in outputs] == [71] + assert speculator.sessions == [] + + +def test_verify_admission_and_diagnostics_report_temporary_dynamic_mode() -> ( + None +): + speculator = _Speculator() + engine = _engine( + speculator, + verify_token_budget=2, + verify_expert_byte_budget=32, + ) + engine.add_request( + "admit", + [60], + SamplingParams(temperature=0.5, max_tokens=3), + ) + assert [output.token_id for output in engine.step()] == [61] + + assert engine.step() == [] + record = next(iter(engine.speculative_sessions.values())) + assert record.pending_draft is not None + assert engine.get_stats()["speculative_execution_context"] == ( + EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + ) + + outputs = engine.step() + assert [output.token_id for output in outputs] == [62, 63] + assert outputs[-1].finished + assert record.decode_state.invariant_holds() + assert record.decode_state.cached_len == record.decode_state.prompt_len + 3 diff --git a/tests/python/serving/test_dflash_stage4b.py b/tests/python/serving/test_dflash_stage4b.py new file mode 100644 index 00000000..1f375fa2 --- /dev/null +++ b/tests/python/serving/test_dflash_stage4b.py @@ -0,0 +1,856 @@ +from __future__ import annotations + +import types +from dataclasses import dataclass, field + +import pytest +import torch + +from moe_infinity.serving.engine import ContinuousBatchingEngine +from moe_infinity.serving.mla_cache import MLAPagedKVCache +from moe_infinity.serving.sequence import SamplingParams +from moe_infinity.serving.spec_cache_adapter import ( + EXECUTION_CONTEXT_PAGED_MLA, + PagedCacheAdapter, + PagedCacheSnapshot, +) +from moe_infinity.serving.spec_session_driver import ( + EXECUTION_CONTEXT_TEMPORARY_DYNAMIC, + SpecSessionDriver, + TemporaryDynamicCacheContext, +) +from moe_infinity.spec_decode.dflash import DFlashSpeculator +from moe_infinity.spec_decode.protocols import RichForwardResult + + +def _cache() -> MLAPagedKVCache: + return MLAPagedKVCache( + num_blocks=8, + block_size=2, + num_layers=2, + latent_dim=3, + rope_dim=2, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + +def test_adapter_append_snapshot_truncate_and_metadata_use_owned_pages() -> ( + None +): + cache = _cache() + adapter = PagedCacheAdapter(cache, seq_id=7, initial_length=3) + original_pages = tuple(cache.get_block_table(7)) + + snapshot = adapter.snapshot() + adapter.append(4) + metadata = adapter.build_attention_metadata( + query_length=4, is_prefill=False + ) + + assert snapshot == PagedCacheSnapshot(3, original_pages) + assert adapter.cache_kind == "paged" + assert adapter.mode == EXECUTION_CONTEXT_PAGED_MLA + assert metadata.seq_id == 7 + assert metadata.seq_lens.tolist() == [7] + assert metadata.block_tables[0, :4].tolist() == cache.get_block_table(7) + assert metadata.slot_mapping.tolist() == [ + cache.get_block_table(7)[1] * 2 + 1, + cache.get_block_table(7)[2] * 2, + cache.get_block_table(7)[2] * 2 + 1, + cache.get_block_table(7)[3] * 2, + ] + + adapter.truncate(5) + assert adapter.logical_length() == 5 + adapter.restore(snapshot) + assert adapter.logical_length() == 3 + assert tuple(cache.get_block_table(7)) == original_pages + + +def test_adapter_isolates_sequences_and_release_frees_only_its_owner() -> None: + cache = _cache() + first = PagedCacheAdapter(cache, seq_id=1, initial_length=2) + second = PagedCacheAdapter(cache, seq_id=2, initial_length=2) + second_pages = tuple(cache.get_block_table(2)) + + first.append(3) + first.truncate(1) + + assert tuple(cache.get_block_table(2)) == second_pages + assert second.logical_length() == 2 + first.release() + with pytest.raises(KeyError, match="unknown sequence id: 1"): + cache.get_block_table(1) + assert tuple(cache.get_block_table(2)) == second_pages + + +def test_adapter_release_is_idempotent_and_rejects_later_mutation() -> None: + adapter = PagedCacheAdapter(_cache(), seq_id=3, initial_length=1) + + adapter.release() + adapter.release() + + with pytest.raises(RuntimeError, match="released"): + adapter.append(1) + with pytest.raises(RuntimeError, match="released"): + adapter.snapshot() + + +def test_adapter_reports_resident_only_preemption_without_duplicate_storage() -> ( + None +): + cache = _cache() + adapter = PagedCacheAdapter(cache, seq_id=4, initial_length=3) + storage = cache.get_mla_cache_tensors() + pages = tuple(cache.get_block_table(4)) + + assert adapter.swap_out() is False + assert adapter.swap_in() is False + assert adapter.cache is cache + assert adapter.cache.get_mla_cache_tensors() is storage + assert tuple(cache.get_block_table(4)) == pages + + +def test_dflash_target_forward_passes_mla_metadata_and_returns_engine_handle() -> ( + None +): + cache = _cache() + adapter = PagedCacheAdapter(cache, seq_id=5, initial_length=2) + metadata = adapter.build_attention_metadata(query_length=2, is_prefill=True) + calls: list[object] = [] + + def rich( + token_ids: list[int], + attention_metadata: object, + logits_to_keep: int = 0, + ) -> RichForwardResult: + calls.append(attention_metadata) + return RichForwardResult( + logits=torch.zeros(1, len(token_ids), 8), + hidden_states=(torch.zeros(1, len(token_ids), 4),), + cache_handle=cache, + ) + + speculator = DFlashSpeculator.__new__(DFlashSpeculator) + speculator.moe = types.SimpleNamespace(_native_model_forward_rich=rich) + + _, _, handle = speculator._forward_target( + torch.tensor([[1, 2]]), + past_key_values=adapter, + logits_to_keep=1, + attention_metadata=metadata, + ) + + assert calls == [metadata] + assert handle is cache + + +class _DraftCache: + def __init__(self) -> None: + self.released = False + + def crop(self, length: int) -> None: + if length == 0: + self.released = True + + +@dataclass +class _DriverSession: + emitted: list[int] + target_kv: object + draft_kv: _DraftCache = field(default_factory=_DraftCache) + finished: bool = False + + def clear_pending(self) -> None: + pass + + +class _PagedDriverSpeculator: + def __init__( + self, cache: MLAPagedKVCache, *, dflash_block_size: int = 2 + ) -> None: + self.moe = types.SimpleNamespace( + _native_mla_cache=cache, + _cached_past_key_values=object(), + _get_mla_attention_modules=lambda: [object()], + ) + self.received_adapter: object | None = None + self.begin_calls: list[dict[str, object]] = [] + self.begin_error: BaseException | None = None + self.append_on_draft = 0 + self.config = types.SimpleNamespace(block_size=dflash_block_size) + + def begin_session( + self, + input_ids: torch.Tensor, + *, + target_cache_adapter: object | None = None, + **kwargs: object, + ) -> _DriverSession: + self.begin_calls.append( + {**kwargs, "target_cache_adapter": target_cache_adapter} + ) + self.received_adapter = target_cache_adapter + if self.begin_error is not None: + raise self.begin_error + return _DriverSession( + emitted=[int(input_ids[0, -1]) + 1], + target_kv=( + target_cache_adapter + if target_cache_adapter is not None + else object() + ), + ) + + def draft_round(self, session: _DriverSession) -> object: + if self.append_on_draft: + session.target_kv.append(self.append_on_draft) + return types.SimpleNamespace(tokens=2, expert_bytes=0) + + def verify_round(self, session: _DriverSession) -> object: + del session + return types.SimpleNamespace(committed_count=1) + + +def test_standalone_driver_keeps_paged_mla_default_off() -> None: + driver = SpecSessionDriver(_PagedDriverSpeculator(_cache())) + + record = _begin_driver_record(driver, "standalone-off", 5, [1, 2]) + + assert record.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert driver.admission_stats["enabled"] is False + + +@pytest.mark.parametrize( + "block_size", + [pytest.param(None, id="missing"), pytest.param("5", id="non-int"), 0, 1], +) +def test_invalid_dflash_block_size_falls_back_without_mla_allocation( + block_size: object, +) -> None: + cache = _cache() + speculator = _PagedDriverSpeculator(cache) + if block_size is None: + del speculator.config.block_size + else: + speculator.config.block_size = block_size + driver = SpecSessionDriver(speculator, enable_paged_mla=True) + + record = _begin_driver_record(driver, "invalid-block", 48, [1, 2]) + + assert record.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert record.diagnostics()["paged_mla_admission"] == { + "eligible": False, + "admitted": False, + "reason": "ineligible", + } + assert speculator.received_adapter is None + assert cache.free_block_count == cache.num_blocks + + +def test_driver_selects_paged_mla_for_eligible_greedy_session_without_dynamic_context() -> ( + None +): + speculator = _PagedDriverSpeculator(_cache()) + previous_dense_owner = speculator.moe._cached_past_key_values + driver = SpecSessionDriver(speculator, enable_paged_mla=True) + + record = driver.begin( + request_id="paged", + seq_id=6, + prompt_token_ids=[1, 2], + max_new_tokens=3, + temperature=0.0, + top_k=0, + top_p=1.0, + stop_token_ids=(), + callbacks=(), + ) + + assert record.execution_context.mode == EXECUTION_CONTEXT_PAGED_MLA + assert driver.execution_context_mode == EXECUTION_CONTEXT_PAGED_MLA + assert not isinstance( + record.execution_context, TemporaryDynamicCacheContext + ) + assert record.spec_session.target_kv is speculator.received_adapter + assert isinstance(record.spec_session.target_kv, PagedCacheAdapter) + assert ( + record.spec_session.target_kv.cache is speculator.moe._native_mla_cache + ) + assert speculator.moe._cached_past_key_values is previous_dense_owner + + draft_cache = record.spec_session.draft_kv + driver.release(record) + assert draft_cache.released + with pytest.raises(KeyError, match="unknown sequence id: 6"): + speculator.moe._native_mla_cache.get_block_table(6) + + +def test_driver_cancel_releases_paged_target_and_drafter_owners() -> None: + speculator = _PagedDriverSpeculator(_cache()) + driver = SpecSessionDriver(speculator, enable_paged_mla=True) + record = driver.begin( + request_id="cancel", + seq_id=12, + prompt_token_ids=[1, 2], + max_new_tokens=3, + temperature=0.0, + top_k=0, + top_p=1.0, + stop_token_ids=(), + callbacks=(), + ) + draft_cache = record.spec_session.draft_kv + + driver.cancel(12) + + assert record.cancelled and record.released + assert draft_cache.released + assert 12 not in driver.sessions + with pytest.raises(KeyError, match="unknown sequence id: 12"): + speculator.moe._native_mla_cache.get_block_table(12) + + +def test_driver_keeps_sampled_request_on_explicit_stage4a_fallback() -> None: + speculator = _PagedDriverSpeculator(_cache()) + driver = SpecSessionDriver(speculator, enable_paged_mla=True) + + record = driver.begin( + request_id="sampled", + seq_id=7, + prompt_token_ids=[1], + max_new_tokens=2, + temperature=0.7, + top_k=0, + top_p=1.0, + stop_token_ids=(), + callbacks=(), + ) + + assert record.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert isinstance(record.execution_context, TemporaryDynamicCacheContext) + assert speculator.received_adapter is None + assert speculator.begin_calls[0]["temperature"] == 0.7 + + +def test_driver_caps_concurrent_resident_paged_mla_sessions() -> None: + speculator = _PagedDriverSpeculator(_cache()) + driver = SpecSessionDriver( + speculator, + enable_paged_mla=True, + max_resident_paged_speculative_sessions=1, + min_free_mla_blocks_after_admission=1, + ) + + first = _begin_driver_record(driver, "first", 31, [1, 2]) + second = _begin_driver_record(driver, "second", 32, [3, 4]) + + assert first.execution_context.mode == EXECUTION_CONTEXT_PAGED_MLA + assert second.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert second.diagnostics()["paged_mla_admission"] == { + "eligible": True, + "admitted": False, + "reason": "session_cap", + } + assert driver.admission_stats["counters"]["session_cap"] == 1 + + +def test_driver_preserves_free_block_reserve_with_immediate_stage4a_fallback() -> ( + None +): + cache = _cache() + speculator = _PagedDriverSpeculator(cache) + driver = SpecSessionDriver( + speculator, + enable_paged_mla=True, + max_resident_paged_speculative_sessions=2, + min_free_mla_blocks_after_admission=8, + ) + + record = _begin_driver_record(driver, "reserve", 33, [1, 2]) + + assert record.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert record.diagnostics()["paged_mla_admission"]["reason"] == ( + "free_block_reserve" + ) + assert cache.free_block_count == 8 + assert driver.admission_stats["counters"]["free_block_reserve"] == 1 + + +def test_driver_admits_when_transient_peak_plus_reserve_exactly_fits() -> None: + cache = MLAPagedKVCache(6, 2, 1, 3, 2, torch.float32, torch.device("cpu")) + driver = SpecSessionDriver( + _PagedDriverSpeculator(cache, dflash_block_size=5), + enable_paged_mla=True, + max_resident_paged_speculative_sessions=1, + min_free_mla_blocks_after_admission=1, + ) + + record = _begin_driver_record(driver, "exact", 40, [1, 2], max_new_tokens=4) + + assert record.execution_context.mode == EXECUTION_CONTEXT_PAGED_MLA + assert record.diagnostics()["paged_mla_admission"]["reason"] == "admitted" + + +def test_driver_falls_back_when_transient_peak_is_one_block_short() -> None: + cache = MLAPagedKVCache(5, 2, 1, 3, 2, torch.float32, torch.device("cpu")) + driver = SpecSessionDriver( + _PagedDriverSpeculator(cache, dflash_block_size=5), + enable_paged_mla=True, + max_resident_paged_speculative_sessions=1, + min_free_mla_blocks_after_admission=1, + ) + + record = _begin_driver_record(driver, "short", 41, [1, 2], max_new_tokens=4) + + assert record.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert record.diagnostics()["paged_mla_admission"]["reason"] == ( + "free_block_reserve" + ) + assert cache.free_block_count == 5 + + +def test_reserved_transient_peak_prevents_mid_verify_exhaustion() -> None: + cache = MLAPagedKVCache(6, 2, 1, 3, 2, torch.float32, torch.device("cpu")) + driver = SpecSessionDriver( + _PagedDriverSpeculator(cache, dflash_block_size=5), + enable_paged_mla=True, + max_resident_paged_speculative_sessions=2, + min_free_mla_blocks_after_admission=1, + ) + first = _begin_driver_record( + driver, "peak-owner", 46, [1, 2], max_new_tokens=4 + ) + + competing = _begin_driver_record( + driver, "peak-competitor", 47, [3, 4, 5], max_new_tokens=1 + ) + assert ( + competing.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + ) + + first.spec_session.target_kv.append(8) + assert len(cache.get_block_table(46)) == 5 + assert cache.free_block_count == 1 + + +def test_driver_accounts_active_declared_headroom_and_release() -> None: + cache = MLAPagedKVCache(8, 2, 1, 3, 2, torch.float32, torch.device("cpu")) + driver = SpecSessionDriver( + _PagedDriverSpeculator(cache, dflash_block_size=5), + enable_paged_mla=True, + max_resident_paged_speculative_sessions=2, + min_free_mla_blocks_after_admission=1, + ) + first = _begin_driver_record( + driver, "first-budget", 42, [1, 2], max_new_tokens=4 + ) + + blocked = _begin_driver_record( + driver, "blocked-budget", 43, [1, 2], max_new_tokens=2 + ) + assert blocked.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert blocked.diagnostics()["paged_mla_admission"]["reason"] == ( + "free_block_reserve" + ) + + driver.release(first) + admitted = _begin_driver_record( + driver, "released-budget", 44, [1, 2], max_new_tokens=2 + ) + assert admitted.execution_context.mode == EXECUTION_CONTEXT_PAGED_MLA + + +def test_driver_counts_backend_begin_failure_instead_of_admission() -> None: + cache = _cache() + speculator = _PagedDriverSpeculator(cache) + speculator.begin_error = RuntimeError("begin failed") + driver = SpecSessionDriver(speculator, enable_paged_mla=True) + + with pytest.raises(RuntimeError, match="begin failed"): + _begin_driver_record(driver, "begin-failed", 45, [1, 2]) + + assert cache.free_block_count == 8 + assert driver.admission_stats["counters"]["admitted"] == 0 + assert driver.admission_stats["counters"]["begin_failed"] == 1 + + +def test_driver_admits_after_prior_paged_session_releases() -> None: + driver = SpecSessionDriver( + _PagedDriverSpeculator(_cache()), + enable_paged_mla=True, + max_resident_paged_speculative_sessions=1, + min_free_mla_blocks_after_admission=1, + ) + first = _begin_driver_record(driver, "first", 34, [1, 2]) + blocked = _begin_driver_record(driver, "blocked", 35, [1, 2]) + assert blocked.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + + driver.release(first) + admitted = _begin_driver_record(driver, "admitted", 36, [1, 2]) + + assert admitted.execution_context.mode == EXECUTION_CONTEXT_PAGED_MLA + assert admitted.diagnostics()["paged_mla_admission"]["reason"] == "admitted" + assert driver.admission_stats["active_sessions"] == 1 + + +def test_engine_stats_expose_paged_mla_admission_decisions() -> None: + driver = SpecSessionDriver( + _PagedDriverSpeculator(_cache()), + enable_paged_mla=True, + max_resident_paged_speculative_sessions=1, + min_free_mla_blocks_after_admission=1, + ) + _begin_driver_record(driver, "first", 37, [1, 2]) + _begin_driver_record(driver, "blocked", 38, [1, 2]) + engine = ContinuousBatchingEngine.__new__(ContinuousBatchingEngine) + engine._spec_session_driver = driver + engine._sequences = {} + engine._request_to_seq_ids = {} + engine._completed_request_ids = set() + engine._cancelled_request_ids = set() + engine._request_failures = {} + engine._num_steps = 0 + engine._total_generated_tokens = 0 + engine.kv_cache = types.SimpleNamespace( + num_blocks=1, + block_allocator=types.SimpleNamespace(num_free_blocks=1), + ) + engine.memory_manager = types.SimpleNamespace(report=lambda: {}) + + stats = engine.get_stats() + + assert stats["paged_mla_admission"] == driver.admission_stats + assert stats["paged_mla_admission"]["counters"]["session_cap"] == 1 + + +class _ServingModel: + config = types.SimpleNamespace(vocab_size=128, eos_token_id=99) + + def eval(self) -> None: + pass + + def forward( + self, input_ids: torch.Tensor, **kwargs: object + ) -> types.SimpleNamespace: + del kwargs + logits = torch.full((*input_ids.shape, 128), -1e9) + for row in range(input_ids.shape[0]): + for col in range(input_ids.shape[1]): + logits[row, col, int(input_ids[row, col]) + 1] = 0 + return types.SimpleNamespace(logits=logits) + + +class _ServingOffload: + def __init__(self) -> None: + self.request_id = 0 + self.expert_tracer = types.SimpleNamespace(create_entry=lambda: 0) + self.expert_layer_modules = [types.SimpleNamespace(seq_id_list=[])] + + def _generate_request_id(self) -> int: + value = self.request_id + self.request_id += 1 + return value + + +def _paged_engine( + speculator: _PagedDriverSpeculator, *, enabled: bool +) -> ContinuousBatchingEngine: + return ContinuousBatchingEngine( + model=_ServingModel(), + engine=_ServingOffload(), + config={ + "device_memory_ratio": 0.75, + "kv_cache_ratio": 0.25, + "max_batch_size": 8, + "max_tokens_per_step": 16, + "block_size": 4, + "num_layers": 1, + "num_kv_heads": 2, + "head_dim": 8, + "dtype": "float32", + "eos_token_id": 99, + "num_kv_blocks": 32, + "verify_token_budget": 8, + "verify_expert_byte_budget": 128, + "verify_token_deficit_cap": 32, + "verify_expert_byte_deficit_cap": 512, + "enable_deepseek_mla_paging": enabled, + "max_resident_paged_speculative_sessions": 1, + "min_free_mla_blocks_after_admission": 1, + }, + speculative_draft=speculator, + ) + + +def test_engine_add_request_step_selects_paged_mla_only_when_enabled() -> None: + enabled_engine = _paged_engine( + _PagedDriverSpeculator(_cache()), enabled=True + ) + enabled_engine.add_request( + "enabled", [1, 2], SamplingParams(temperature=0.0, max_tokens=3) + ) + enabled_engine.step() + enabled_record = next(iter(enabled_engine.speculative_sessions.values())) + assert enabled_record.execution_context.mode == EXECUTION_CONTEXT_PAGED_MLA + + default_off_engine = _paged_engine( + _PagedDriverSpeculator(_cache()), enabled=False + ) + default_off_engine.add_request( + "default-off", [1, 2], SamplingParams(temperature=0.0, max_tokens=3) + ) + default_off_engine.step() + default_record = next( + iter(default_off_engine.speculative_sessions.values()) + ) + assert ( + default_record.execution_context.mode + == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + ) + + +def test_engine_append_exhaustion_records_failure_and_releases_paged_owner() -> ( + None +): + cache = MLAPagedKVCache(5, 2, 1, 3, 2, torch.float32, torch.device("cpu")) + speculator = _PagedDriverSpeculator(cache) + speculator.append_on_draft = 2 + engine = _paged_engine(speculator, enabled=True) + engine.add_request( + "exhaust", [1, 2], SamplingParams(temperature=0.0, max_tokens=4) + ) + engine.step() + seq_id = engine._request_to_seq_ids["exhaust"][0] + cache.allocate_sequence(999, 8) + + with pytest.raises(RuntimeError, match="BlockAllocator exhausted"): + engine.step() + + assert engine.get_request_failure("exhaust") == { + "phase": "draft", + "failure_type": "RuntimeError", + "code": "speculative_draft_failed", + } + assert engine.speculative_sessions == {} + assert engine.has_pending_requests() is False + with pytest.raises(KeyError, match=f"unknown sequence id: {seq_id}"): + cache.get_block_table(seq_id) + + +def _begin_driver_record( + driver: SpecSessionDriver, + request_id: str, + seq_id: int, + prompt_token_ids: list[int], + max_new_tokens: int = 3, +): + return driver.begin( + request_id=request_id, + seq_id=seq_id, + prompt_token_ids=prompt_token_ids, + max_new_tokens=max_new_tokens, + temperature=0.0, + top_k=0, + top_p=1.0, + stop_token_ids=(), + callbacks=(), + ) + + +def test_driver_keeps_non_deepseek_owner_on_explicit_stage4a_fallback() -> None: + speculator = _PagedDriverSpeculator(_cache()) + speculator.moe._get_mla_attention_modules = lambda: [] + driver = SpecSessionDriver(speculator, enable_paged_mla=True) + + record = driver.begin( + request_id="qwen-or-hybrid", + seq_id=8, + prompt_token_ids=[1], + max_new_tokens=2, + temperature=0.0, + top_k=0, + top_p=1.0, + stop_token_ids=(), + callbacks=(), + ) + + assert record.execution_context.mode == EXECUTION_CONTEXT_TEMPORARY_DYNAMIC + assert speculator.received_adapter is None + + +def test_canonical_session_prefill_and_verify_use_transient_mla_slots_then_truncate() -> ( + None +): + cache = _cache() + adapter = PagedCacheAdapter(cache, seq_id=10, initial_length=2) + speculator = DFlashSpeculator.__new__(DFlashSpeculator) + speculator.device = "cpu" + speculator.config = types.SimpleNamespace( + block_size=2, + target_layer_ids=[0], + mask_token_id=0, + ) + speculator._drafter_has_kv_cache = False + speculator.moe = types.SimpleNamespace(engine=None) + speculator.target = types.SimpleNamespace( + config=types.SimpleNamespace(eos_token_id=None), + generation_config=types.SimpleNamespace(eos_token_id=None), + ) + speculator._configure_target_hooks = lambda input_ids: None + speculator.route_ahead_stats = None + metadata_seen: list[object] = [] + + def forward_target( + input_ids: torch.Tensor, + past_key_values: object = None, + logits_to_keep: int = 0, + *, + attention_metadata: object = None, + **kwargs: object, + ) -> tuple[torch.Tensor, tuple[torch.Tensor, ...], object]: + del logits_to_keep, kwargs + metadata_seen.append(attention_metadata) + query_len = int(input_ids.shape[1]) + logits = torch.zeros(1, query_len, 16) + if query_len == 2 and past_key_values is adapter: + logits[0, 0, 9] = 1 + logits[0, 1, 8] = 1 + else: + logits[0, -1, 3] = 1 + hidden = ( + torch.zeros(1, query_len, 4), + torch.zeros(1, query_len, 4), + ) + return logits, hidden, cache + + speculator._forward_target = forward_target + session = speculator.begin_session( + torch.tensor([[1, 2]]), + max_new_tokens=2, + temperature=0.0, + stop_token_ids=[], + target_cache_adapter=adapter, + ) + + assert session.target_kv is adapter + assert metadata_seen[0].is_prefill is True + assert metadata_seen[0].seq_id == 10 + assert adapter.logical_length() == 2 + + block = torch.tensor([[session.anchor, 4]]) + session._pending = True + session._pending_block = block + session._pending_prev_start = 2 + session._pending_cache_snapshot = adapter.snapshot() + session.pending_draft_probs = None + + result = speculator.verify_round(session) + + assert result.accept == 0 + assert metadata_seen[1].is_prefill is False + assert metadata_seen[1].seq_lens.tolist() == [4] + assert metadata_seen[1].slot_mapping.numel() == 2 + assert adapter.logical_length() == 3 + assert session.start == 3 + assert session.step_trace[-1].target_cache_len == 3 + + +def test_real_deepseek_rich_forward_writes_the_adapter_owned_cache_for_verify() -> ( + None +): + transformers = pytest.importorskip("transformers") + pytest.importorskip("transformers.models.deepseek_v3.modeling_deepseek_v3") + from moe_infinity.entrypoints.big_modeling import MoE + from moe_infinity.models.deepseek_mla_attention import adapt_deepseek_model + + config = transformers.DeepseekV3Config( + hidden_size=16, + intermediate_size=24, + num_hidden_layers=1, + num_attention_heads=2, + num_key_value_heads=2, + q_lora_rank=None, + kv_lora_rank=4, + qk_nope_head_dim=4, + qk_rope_head_dim=4, + v_head_dim=4, + n_routed_experts=2, + n_shared_experts=1, + num_experts_per_tok=1, + moe_intermediate_size=8, + n_group=1, + topk_group=1, + rope_interleave=True, + vocab_size=32, + first_k_dense_replace=1, + ) + model = transformers.DeepseekV3ForCausalLM(config).eval() + cache = MLAPagedKVCache(4, 2, 1, 4, 4, torch.float32, torch.device("cpu")) + assert len(adapt_deepseek_model(model, cache, enabled=True)) == 1 + adapter = PagedCacheAdapter(cache, seq_id=11, initial_length=2) + + shell = MoE.__new__(MoE) + shell.model = model + shell._cached_past_key_values = None + shell._native_attention_backend = None + shell._native_mla_cache = cache + shell._resolve_native_input_device = lambda: torch.device("cpu") + + prefill = shell._native_model_forward_rich( + [1, 2], + adapter.build_attention_metadata(query_length=2, is_prefill=True), + ) + before_verify = cache.get_mla_cache_tensors().clone() + adapter.append(2) + verify = shell._native_model_forward_rich( + [3, 4], + adapter.build_attention_metadata(query_length=2, is_prefill=False), + ) + + assert isinstance(prefill, RichForwardResult) + assert isinstance(verify, RichForwardResult) + assert prefill.cache_handle is cache and verify.cache_handle is cache + assert shell._cached_past_key_values is None + assert not torch.equal(cache.get_mla_cache_tensors(), before_verify) + + +def test_rich_verify_uses_one_position_per_tentative_token( + monkeypatch: pytest.MonkeyPatch, +) -> None: + import moe_infinity.models.deepseek_mla_attention as mla_attention + from moe_infinity.entrypoints.big_modeling import MoE + + calls: list[dict[str, object]] = [] + + class _Model: + def __call__(self, input_ids: torch.Tensor, **kwargs: object) -> object: + calls.append(dict(kwargs)) + return types.SimpleNamespace( + logits=torch.zeros(1, input_ids.shape[1], 8), + hidden_states=(torch.zeros(1, input_ids.shape[1], 4),), + ) + + monkeypatch.setattr( + mla_attention, "set_deepseek_mla_context", lambda *args: None + ) + monkeypatch.setattr( + mla_attention, "clear_deepseek_mla_context", lambda *args: None + ) + shell = MoE.__new__(MoE) + shell.model = _Model() + shell._native_mla_cache = object() + shell._native_attention_backend = None + shell._resolve_native_input_device = lambda: torch.device("cpu") + shell._get_mla_attention_modules = lambda: [object()] + metadata = types.SimpleNamespace( + block_tables=torch.tensor([[0, 1]], dtype=torch.int32), + seq_lens=torch.tensor([4], dtype=torch.int32), + slot_mapping=torch.tensor([2, 3], dtype=torch.int64), + is_prefill=False, + ) + + shell._native_model_forward_rich([3, 4], metadata) + + assert calls[0]["position_ids"].tolist() == [[2, 3]] diff --git a/tests/python/serving/test_kv_cache.py b/tests/python/serving/test_kv_cache.py index e24ca02b..3a5f0482 100644 --- a/tests/python/serving/test_kv_cache.py +++ b/tests/python/serving/test_kv_cache.py @@ -1,6 +1,7 @@ import importlib.util import sys from pathlib import Path +from types import ModuleType from typing import Protocol, cast import pytest @@ -9,9 +10,8 @@ ROOT = str(Path(__file__).resolve().parents[3]) if ROOT not in sys.path: sys.path.insert(0, ROOT) -_ = sys.modules.pop("moe_infinity", None) -_ = sys.modules.pop("moe_infinity.serving", None) KV_CACHE_PATH = Path(ROOT) / "moe_infinity" / "serving" / "kv_cache.py" +_MISSING_MODULE = object() class BlockAllocatorProtocol(Protocol): @@ -69,8 +69,15 @@ def _load_classes() -> ( if spec is None or spec.loader is None: raise RuntimeError(f"failed to load module from {KV_CACHE_PATH}") module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) + previous_module = sys.modules.get(module_name, _MISSING_MODULE) + try: + sys.modules[module_name] = module + spec.loader.exec_module(module) + finally: + if previous_module is _MISSING_MODULE: + _ = sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = cast(ModuleType, previous_module) return ( cast(type[BlockAllocatorProtocol], getattr(module, "BlockAllocator")), cast(type[PagedKVCacheProtocol], getattr(module, "PagedKVCache")), diff --git a/tests/python/serving/test_memory_manager.py b/tests/python/serving/test_memory_manager.py index bfe88ba7..2dc8b304 100644 --- a/tests/python/serving/test_memory_manager.py +++ b/tests/python/serving/test_memory_manager.py @@ -2,6 +2,7 @@ import json import sys from pathlib import Path +from types import ModuleType from typing import Optional, Protocol, Union, cast import torch @@ -9,11 +10,10 @@ ROOT = str(Path(__file__).resolve().parents[3]) if ROOT not in sys.path: sys.path.insert(0, ROOT) -_ = sys.modules.pop("moe_infinity", None) -_ = sys.modules.pop("moe_infinity.serving", None) MEMORY_MANAGER_PATH = ( Path(ROOT) / "moe_infinity" / "serving" / "memory_manager.py" ) +_MISSING_MODULE = object() class MemoryBudgetProtocol(Protocol): @@ -77,8 +77,15 @@ def _load_classes() -> ( if spec is None or spec.loader is None: raise RuntimeError(f"failed to load module from {MEMORY_MANAGER_PATH}") module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) + previous_module = sys.modules.get(module_name, _MISSING_MODULE) + try: + sys.modules[module_name] = module + spec.loader.exec_module(module) + finally: + if previous_module is _MISSING_MODULE: + _ = sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = cast(ModuleType, previous_module) return ( cast(type[MemoryBudgetProtocol], getattr(module, "MemoryBudget")), cast(type[MemoryManagerProtocol], getattr(module, "MemoryManager")), diff --git a/tests/python/serving/test_mla_paged_cache.py b/tests/python/serving/test_mla_paged_cache.py new file mode 100644 index 00000000..e5d1e267 --- /dev/null +++ b/tests/python/serving/test_mla_paged_cache.py @@ -0,0 +1,131 @@ +from __future__ import annotations + +import pytest +import torch + +from moe_infinity.serving.mla_cache import MLAPagedKVCache + + +def _cache() -> MLAPagedKVCache: + return MLAPagedKVCache( + num_blocks=4, + block_size=2, + num_layers=2, + latent_dim=3, + rope_dim=2, + dtype=torch.float32, + device=torch.device("cpu"), + ) + + +def test_mla_cache_writes_external_slots_without_cross_layer_aliasing() -> None: + cache = _cache() + cache.allocate_sequence(7, 3) + block_table = cache.get_block_table(7) + slots = torch.tensor( + [block_table[0] * 2, block_table[1] * 2], dtype=torch.int64 + ) + latent = torch.tensor([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]) + rope = torch.tensor([[7.0, 8.0], [9.0, 10.0]]) + + cache.write(1, latent, rope, slots) + + packed = cache.get_mla_cache_tensors() + assert packed.shape == (2, 4, 2, 5) + assert torch.equal( + packed[1, block_table[0], 0], torch.tensor([1, 2, 3, 7, 8]) + ) + assert torch.equal( + packed[1, block_table[1], 0], torch.tensor([4, 5, 6, 9, 10]) + ) + assert torch.count_nonzero(packed[0]) == 0 + + +def test_mla_cache_reads_block_table_order_and_truncates_logically() -> None: + cache = _cache() + cache.allocate_sequence(9, 4) + block_table = cache.get_block_table(9) + slots = torch.tensor( + [ + block_table[0] * 2, + block_table[0] * 2 + 1, + block_table[1] * 2, + block_table[1] * 2 + 1, + ] + ) + latent = torch.arange(12, dtype=torch.float32).view(4, 3) + rope = torch.arange(8, dtype=torch.float32).view(4, 2) + 20 + cache.write(0, latent, rope, slots) + + read_latent, read_rope = cache.read(0, block_table, seq_len=4) + assert torch.equal(read_latent, latent) + assert torch.equal(read_rope, rope) + + cache.truncate_tokens(9, 2) + assert cache.get_block_table(9) == block_table[:1] + assert cache.block_allocator.num_free_blocks == 3 + + +def test_mla_cache_rejects_invalid_layer_and_slot() -> None: + cache = _cache() + with pytest.raises(IndexError, match="layer_idx"): + cache.write(2, torch.zeros(1, 3), torch.zeros(1, 2), torch.tensor([0])) + with pytest.raises(ValueError, match="past allocated pages"): + cache.write(0, torch.zeros(1, 3), torch.zeros(1, 2), torch.tensor([8])) + + +def test_mla_cache_vectorized_write_matches_reference_with_duplicate_slots( + monkeypatch: pytest.MonkeyPatch, +) -> None: + cache = _cache() + slots = torch.tensor([3, 0, 3, 6], dtype=torch.int64) + latent = torch.arange(12, dtype=torch.float32).view(4, 3) + rope = torch.arange(8, dtype=torch.float32).view(4, 2) + 30 + packed = torch.cat((latent, rope), dim=-1) + expected = torch.zeros_like(cache.get_mla_cache_tensors()) + for token_idx, slot in enumerate((3, 0, 3, 6)): + page, offset = divmod(slot, cache.block_size) + expected[1, page, offset] = packed[token_idx] + + original_tolist = torch.Tensor.tolist + + def reject_tolist(tensor: torch.Tensor): + if tensor is slots: + raise AssertionError( + "MLA write must not copy slot_mapping to a Python list" + ) + return original_tolist(tensor) + + monkeypatch.setattr(torch.Tensor, "tolist", reject_tolist) + cache.write(1, latent, rope, slots) + + assert torch.equal(cache.get_mla_cache_tensors(), expected) + assert torch.count_nonzero(cache.get_mla_cache_tensors()[0]) == 0 + + +@pytest.mark.parametrize( + ("slot", "message"), [(-1, "negative"), (8, "past allocated")] +) +def test_mla_cache_vectorized_write_preserves_bounds_errors( + slot: int, message: str +) -> None: + cache = _cache() + with pytest.raises(ValueError, match=message): + cache.write( + 0, + torch.zeros(1, 3), + torch.zeros(1, 2), + torch.tensor([slot]), + ) + + +def test_mla_cache_exposes_free_blocks_across_sequence_lifecycle() -> None: + cache = _cache() + + assert cache.free_block_count == 4 + cache.allocate_sequence(21, 3) + assert cache.free_block_count == 2 + cache.append_tokens(21, 2) + assert cache.free_block_count == 1 + cache.free_sequence(21) + assert cache.free_block_count == 4 diff --git a/tests/python/serving/test_prefix_cache.py b/tests/python/serving/test_prefix_cache.py index 7c102ec2..1d279a80 100644 --- a/tests/python/serving/test_prefix_cache.py +++ b/tests/python/serving/test_prefix_cache.py @@ -1,14 +1,14 @@ import importlib.util import sys from pathlib import Path +from types import ModuleType from typing import Callable, Protocol, cast ROOT = str(Path(__file__).resolve().parents[3]) if ROOT not in sys.path: sys.path.insert(0, ROOT) -_ = sys.modules.pop("moe_infinity", None) -_ = sys.modules.pop("moe_infinity.serving", None) PREFIX_CACHE_PATH = Path(ROOT) / "moe_infinity" / "serving" / "prefix_cache.py" +_MISSING_MODULE = object() class PrefixCacheProtocol(Protocol): @@ -48,8 +48,15 @@ def _load_prefix_cache_objects() -> ( raise RuntimeError(f"failed to load module from {PREFIX_CACHE_PATH}") module = importlib.util.module_from_spec(spec) - sys.modules[module_name] = module - spec.loader.exec_module(module) + previous_module = sys.modules.get(module_name, _MISSING_MODULE) + try: + sys.modules[module_name] = module + spec.loader.exec_module(module) + finally: + if previous_module is _MISSING_MODULE: + _ = sys.modules.pop(module_name, None) + else: + sys.modules[module_name] = cast(ModuleType, previous_module) return ( cast(type[PrefixCacheProtocol], getattr(module, "PrefixCache")), cast(Callable[[list[int]], str], getattr(module, "hash_token_block")), diff --git a/tests/python/serving/test_rich_batch_runner.py b/tests/python/serving/test_rich_batch_runner.py new file mode 100644 index 00000000..9f8e0b9e --- /dev/null +++ b/tests/python/serving/test_rich_batch_runner.py @@ -0,0 +1,101 @@ +from __future__ import annotations + +import types + +import torch + +from moe_infinity.serving.batch import BatchMetadata +from moe_infinity.serving.model_runner import ModelRunner +from moe_infinity.serving.sequence import SamplingParams +from moe_infinity.spec_decode.protocols import RichForwardResult + + +class _RichModel: + def __init__(self) -> None: + self.calls = 0 + self.config = types.SimpleNamespace(vocab_size=2) + + def eval(self) -> None: + pass + + def forward(self, input_ids: torch.Tensor, **kwargs: object) -> object: + self.calls += 1 + hidden = input_ids.to(torch.float32).unsqueeze(-1) + return types.SimpleNamespace( + logits=torch.cat((hidden, hidden + 1), dim=-1), + hidden_states=(hidden, hidden + 10), + past_key_values=kwargs.get("past_key_values", "cache"), + ) + + def modules(self) -> list[object]: + return [] + + +class _Engine: + request_id = 0 + expert_layer_modules: list[object] = [] + expert_tracer = None + + +def _batch() -> BatchMetadata: + return BatchMetadata( + seq_ids=[41, 42], + input_token_ids=[3, 4, 9], + seq_lengths=[2, 1], + context_lengths=[0, 7], + is_prefill=[True, False], + block_tables=[[2], [5]], + token_offsets=[0, 2, 3], + sampling_params=[SamplingParams(), SamplingParams()], + ) + + +def test_model_runner_rich_execute_returns_packed_hidden_and_row_cache_handles() -> ( + None +): + model = _RichModel() + runner = ModelRunner(model, _Engine(), device=torch.device("cpu")) + handles = (object(), object()) + + result = runner.execute_rich(_batch(), cache_handles=handles) + + assert isinstance(result, RichForwardResult) + assert model.calls == 1 + assert result.logits.shape == (3, 2) + assert result.hidden_states[0].shape == (3, 1) + assert result.cache_handles == handles + assert result.row_offsets == (0, 2, 3) + assert result.row_lengths == (2, 1) + + +def test_standard_model_runner_api_remains_logits_only() -> None: + runner = ModelRunner(_RichModel(), _Engine(), device=torch.device("cpu")) + + logits = runner.execute(_batch()) + + assert isinstance(logits, torch.Tensor) + assert logits.shape == (3, 2) + + +def test_empty_rich_runner_result_keeps_row_alignment_without_forward() -> None: + model = _RichModel() + runner = ModelRunner(model, _Engine(), device=torch.device("cpu")) + batch = BatchMetadata( + seq_ids=[9], + input_token_ids=[], + seq_lengths=[0], + context_lengths=[4], + is_prefill=[False], + block_tables=[[1]], + token_offsets=[0, 0], + sampling_params=[SamplingParams()], + ) + + result = runner.execute_rich(batch, cache_handles=("paged-9",)) + + assert isinstance(result, RichForwardResult) + assert result.logits.shape[0] == 0 + assert result.hidden_states == () + assert result.cache_handles == ("paged-9",) + assert result.row_offsets == (0, 0) + assert model.calls == 0 diff --git a/tests/python/serving/test_scheduler.py b/tests/python/serving/test_scheduler.py index 792ece17..b083cf6d 100644 --- a/tests/python/serving/test_scheduler.py +++ b/tests/python/serving/test_scheduler.py @@ -117,6 +117,64 @@ def test_memory_pressure_preemption() -> None: assert req1.sequences[0].status is SequenceStatus.SWAPPED +def test_preemption_skips_speculative_groups_without_orphaning_them() -> None: + cache = _make_cache(num_blocks=3) + scheduler = Scheduler(cache, max_batch_size=8, max_tokens_per_step=128) + draft = _make_group("draft", 1, 4) + decode = _make_group("decode", 2, 4) + verify = _make_group("verify", 3, 4) + for group in (draft, decode, verify): + scheduler.add_request(group) + _ = scheduler.schedule() + draft.sequences[0].set_status(SequenceStatus.DRAFT) + decode.sequences[0].set_status(SequenceStatus.DECODE) + verify.sequences[0].set_status(SequenceStatus.DRAFT) + verify.sequences[0].set_status(SequenceStatus.VERIFY) + draft_table = cache.get_block_table(1) + verify_table = cache.get_block_table(3) + + newcomer = _make_group("new", 4, 4) + scheduler.add_request(newcomer) + output = scheduler.schedule() + + assert output.preempted_seq_ids == [2] + assert output.prefill_seq_ids == [4] + assert scheduler.get_running_seq_ids() == [1, 3, 4] + assert [group.request_id for group in scheduler._running] == [ + "draft", + "verify", + "new", + ] + assert cache.get_block_table(1) == draft_table + assert cache.get_block_table(3) == verify_table + + +def test_preemption_preserves_all_non_preemptible_running_groups() -> None: + cache = _make_cache(num_blocks=2) + scheduler = Scheduler(cache, max_batch_size=8, max_tokens_per_step=128) + draft = _make_group("draft", 1, 4) + verify = _make_group("verify", 2, 4) + scheduler.add_request(draft) + scheduler.add_request(verify) + _ = scheduler.schedule() + draft.sequences[0].set_status(SequenceStatus.DRAFT) + verify.sequences[0].set_status(SequenceStatus.DRAFT) + verify.sequences[0].set_status(SequenceStatus.VERIFY) + + scheduler.add_request(_make_group("blocked", 3, 4)) + output = scheduler.schedule() + + assert output.preempted_seq_ids == [] + assert output.prefill_seq_ids == [] + assert [group.request_id for group in scheduler._running] == [ + "draft", + "verify", + ] + assert scheduler.num_waiting == 1 + assert cache.get_block_table(1) + assert cache.get_block_table(2) + + def test_abort_request() -> None: cache = _make_cache() scheduler = Scheduler(cache, max_batch_size=8, max_tokens_per_step=128) diff --git a/tests/python/unit/test_ci_dflash_validation.py b/tests/python/unit/test_ci_dflash_validation.py new file mode 100644 index 00000000..f2505625 --- /dev/null +++ b/tests/python/unit/test_ci_dflash_validation.py @@ -0,0 +1,24 @@ +from pathlib import Path + +ROOT = Path(__file__).resolve().parents[3] +PR_WORKFLOW = ROOT / ".github/workflows/ci-pr.yml" + + +def test_python_matrix_runs_dflash_unified_validation_e2e() -> None: + workflow = PR_WORKFLOW.read_text(encoding="utf-8") + unit_job = workflow.split(" unit-tests:\n", maxsplit=1)[1].split( + "\n build:", maxsplit=1 + )[0] + + unit_test_position = unit_job.index("- name: Run CPU unit tests") + dflash_position = unit_job.index( + "- name: Run DFlash unified validation E2E" + ) + contextpilot_position = unit_job.index("- name: Install contextpilot") + + assert 'python-version: ["3.10", "3.12"]' in unit_job + assert ( + "pytest tests/python/integration/" + "test_dflash_unified_validation_e2e.py" in unit_job + ) + assert unit_test_position < dflash_position < contextpilot_position diff --git a/tests/python/unit/test_utils_config.py b/tests/python/unit/test_utils_config.py index d2768acf..5a48551c 100644 --- a/tests/python/unit/test_utils_config.py +++ b/tests/python/unit/test_utils_config.py @@ -89,3 +89,30 @@ def test_native_engine_autocorrects_kv_cache_ratio(monkeypatch): kv_cache_memory_ratio=0.0, ) assert config.kv_cache_memory_ratio == pytest.approx(0.15) + + +def test_paged_mla_admission_guard_defaults_are_safe(monkeypatch): + monkeypatch.setattr("torch.cuda.device_count", lambda: 1) + config = ArcherConfig(use_native_engine=False) + + assert config.enable_deepseek_mla_paging is False + assert config.max_resident_paged_speculative_sessions == 1 + assert config.min_free_mla_blocks_after_admission == 1 + + +@pytest.mark.parametrize( + ("field_name", "value"), + [ + ("max_resident_paged_speculative_sessions", True), + ("max_resident_paged_speculative_sessions", -1), + ("min_free_mla_blocks_after_admission", False), + ("min_free_mla_blocks_after_admission", 0), + ], +) +def test_paged_mla_admission_guard_rejects_invalid_values( + monkeypatch, field_name, value +): + monkeypatch.setattr("torch.cuda.device_count", lambda: 1) + + with pytest.raises(ValueError, match=field_name): + ArcherConfig(use_native_engine=False, **{field_name: value}) diff --git a/tests/python/unit/test_watchdog_integration.py b/tests/python/unit/test_watchdog_integration.py index 46d75368..3a3af279 100644 --- a/tests/python/unit/test_watchdog_integration.py +++ b/tests/python/unit/test_watchdog_integration.py @@ -53,9 +53,11 @@ def __init__( class _FakeMoE: + last_config: dict[str, object] | None = None + def __init__(self, model_name: str, config: dict[str, object]) -> None: _ = model_name - _ = config + type(self).last_config = config self.model = SimpleNamespace( config=SimpleNamespace( num_hidden_layers=1, @@ -242,6 +244,80 @@ def test_watchdog_enabled_with_flags(monkeypatch: Any) -> None: _restore_runtime_state(module, original_state) +@pytest.mark.parametrize( + ("mla_overrides", "expected_mla_config"), + [ + ( + {}, + { + "enable_deepseek_mla_paging": False, + "max_resident_paged_speculative_sessions": 1, + "min_free_mla_blocks_after_admission": 1, + }, + ), + ( + { + "enable_deepseek_mla_paging": True, + "max_resident_paged_speculative_sessions": 3, + "min_free_mla_blocks_after_admission": 5, + }, + { + "enable_deepseek_mla_paging": True, + "max_resident_paged_speculative_sessions": 3, + "min_free_mla_blocks_after_admission": 5, + }, + ), + ], +) +def test_partial_namespace_starts_watchdog_and_preserves_mla_config( + monkeypatch: Any, + mla_overrides: dict[str, object], + expected_mla_config: dict[str, object], +) -> None: + module: Any = importlib.import_module(MODULE_NAME) + original_state = _snapshot_runtime_state(module) + _patch_initialize_model_dependencies(monkeypatch, module) + + module.engine = None + module.stream_manager = None + module.tokenizer = None + module.model_name_global = None + module._startup_args = _startup_args( + startup_timeout=12.0, + decode_step_timeout=0.8, + **mla_overrides, + ) + _FakeMoE.last_config = None + + startup_watchdog = MagicMock() + decode_watchdog = MagicMock() + try: + with ( + patch.object( + watchdog_module, + "start_startup_watchdog", + return_value=startup_watchdog, + ) as startup_mock, + patch.object( + watchdog_module, + "start_decode_watchdog", + return_value=decode_watchdog, + ) as decode_mock, + ): + asyncio.run(module._initialize_model()) + + startup_mock.assert_called_once() + decode_mock.assert_called_once() + startup_watchdog.cancel.assert_called_once() + assert isinstance(module.engine, _FakeRuntimeEngine) + assert _FakeMoE.last_config is not None + assert { + key: _FakeMoE.last_config[key] for key in expected_mla_config + } == expected_mla_config + finally: + _restore_runtime_state(module, original_state) + + def test_feed_called_in_engine_loop() -> None: module: Any = importlib.import_module(MODULE_NAME) original_state = _snapshot_runtime_state(module)