model_loader.h

// ORIGINAL packaging helper (no upstream mirror). Factors the model-load +
// engine-stack wiring shared by the OpenAI HTTP server (examples/server/main.cpp)
// and the C ABI (src/capi/vllm_c.cpp) into one place, so both drive the same
// LLMEngine construction path. The wiring itself mirrors the M1.8 LLMEngine
// __init__ (vllm/v1/engine/llm_engine.py @ e24d1b24) and the test harness
// (tests/vllm/v1/test_llm_engine.cpp); this file only owns the pieces + their
// lifetimes so the LLMEngine's by-reference seams stay valid.
#ifndef VLLM_ENTRYPOINTS_MODEL_LOADER_H_
#define VLLM_ENTRYPOINTS_MODEL_LOADER_H_

#include <memory>
#include <mutex>
#include <optional>
#include <string>

#include "vllm/config/kv_transfer.h"
#include "vllm/config/scheduler.h"
#include "vllm/config/speculative.h"
#include "vllm/model_executor/models/model_registry.h"
#include "vllm/model_executor/models/qwen3_5_dense.h"
#include "vllm/model_executor/models/qwen3_5_weights.h"
#include "vllm/model_executor/models/qwen3_dflash.h"  // SPEC-DFLASH D5 draft bundle
#include "vllm/tokenizer/tokenizer.h"
#include "vllm/transformers_utils/hf_config.h"
#include "vllm/v1/core/kv_cache_utils.h"
#include "vllm/v1/core/sched/async_scheduler.h"
#include "vllm/v1/core/sched/scheduler.h"
#include "vllm/v1/engine/core.h"
#include "vllm/v1/engine/async_llm.h"
#include "vllm/v1/engine/input_processor.h"
#include "vllm/v1/engine/llm_engine.h"
#include "vllm/v1/engine/output_processor.h"
#include "vllm/v1/executor/executor.h"
#include "vllm/v1/kv_cache_interface.h"
#include "vllm/v1/kv_offload/kv_connector.h"
#include "vllm/v1/structured_output/manager.h"
#include "vllm/v1/worker/gpu/runner.h"
#include "vt/backend.h"

namespace vllm::entrypoints {

// SPEC-DFLASH D5: the separately-loaded DFlash draft. Unlike MTP (in-target
// mtp.*), the z-lab DFlash draft is its own checkpoint; FromModelDir loads it
// (host-resident bf16 weights + the target-shared embed/lm_head + the resolved
// draft config + k) and the LoadedEngine hands a borrow to the runner. Owned by
// LoadedEngine, declared before runner_ so the borrow outlives it.
struct DflashDraft {
  vllm::Qwen3DFlashWeights weights;
  vllm::HfConfig config;
  int k = 0;
};

// Knobs that size the engine stack. Zero/negative fields fall back to the
// documented defaults (see below), so a default-constructed EngineParams is
// valid.
struct EngineParams {
  int block_size = 32;     // KV block size (tokens/block).
  int num_blocks = 256;    // KV blocks to allocate.
  int max_model_len = 0;   // 0 => config.max_position_embeddings.
  int max_num_seqs = 8;    // max concurrent sequences.
  // Per-step token budget (the chunked-prefill knob). 0 => the bounded PER-ARCH
  // default (see LoadedEngine::ResolveMaxNumBatchedTokens): dense arch 2048 flat
  // (vLLM's DEFAULT_MAX_NUM_BATCHED_TOKENS, vllm/config/scheduler.py:42 @
  // e24d1b24); MoE arch 8192 at max_num_seqs >= 32 else 4096 (GB10-tuned). The
  // budget does NOT scale with max_num_seqs, so a long/many-request prefill is
  // split across steps (enable_chunked_prefill is always true) and the per-step
  // GDN chunked-scan activation stays bounded regardless of concurrency. This is
  // the fix for the 27B 8x1024 conc-8 OOM: the old max_model_len*max_num_seqs
  // product let an 8x1024 (8192-token) prefill run in ONE step, blowing the GDN
  // prefill activation.
  int max_num_batched_tokens = 0;
  // Mirrors vLLM's tri-state --[no-]enable-prefix-caching resolution. Hybrid
  // and attention-free generation models default OFF at the parity pin;
  // decoder-only models default ON. An explicit value overrides the model
  // capability default.
  std::optional<bool> enable_prefix_caching = std::nullopt;
  // Scheduling policy (mirrors SchedulerConfig.policy / the vLLM
  // --scheduling-policy flag). Default fcfs. Set kPriority to schedule by
  // (priority, arrival_time); requests then carry a `priority` field. kLPM is
  // SGLang's cache-aware longest-prefix-match admission ordering
  // (ENG-SGLANG-BEHAVIOR-FLAG SW1) — output-neutral, meaningful only with prefix
  // caching ON (resolves to fcfs otherwise). Exposed on the C ABI as the v9
  // string field vllm_model_params.scheduling_policy = "lpm".
  vllm::SchedulerPolicy policy = vllm::SchedulerPolicy::kFCFS;

  // Jump-forward decoding toggle (ENG-SGLANG-BEHAVIOR-FLAG SW3; the token-unique
  // forced-run subset — see structured_output/jump_forward.h). Tri-state,
  // resolved once at construction by vllm::v1::JumpForwardEnabled(this):
  //   * nullopt (default) => OFF unless the VT_ENABLE_JUMP_FORWARD env override
  //     turns it on (the byte-identical default);
  //   * true  => ON  (unless VT_ENABLE_JUMP_FORWARD is set, which then overrides);
  //   * false => OFF (likewise overridable by the env var).
  // The env var, WHEN SET, always wins (mirrors the VT_ASYNC_SCHED convention).
  // Exposed on the C ABI as vllm_model_params.enable_jump_forward (ABI v10) and
  // on the server as --enable-jump-forward / --disable-jump-forward.
  std::optional<bool> enable_jump_forward = std::nullopt;

  // KV-EXTERNAL-CACHE (LMCache): opt-in external KV-cache connector selection.
  // Empty/absent (default) == NO connector == byte-identical production engine
  // (mirrors vLLM's --kv-transfer-config being unset). When set with a
  // kv_connector name, LoadedEngine builds the connector via KVConnectorFactory,
  // injects the runner's resolved full-attention KV geometry into its
  // extra_config, and wires it to BOTH the scheduler (prefix lookup / prefill
  // shortcut) and the runner (worker-side KV store/load).
  std::optional<vllm::KVTransferConfig> kv_transfer_config = std::nullopt;

  // SPEC-MTP I5d-pre: opt-in speculative-decoding configuration. Empty/absent
  // (default) == NO speculation == byte-identical production engine (mirrors
  // vLLM's --speculative-config being unset). When set to an MTP config,
  // FromModelDir additionally loads the checkpoint's `mtp.*` draft weights
  // (LoadQwen3_5MTP) and retains them on the loaded model so the runner has a
  // typed path to build the draft (LoadedModel::BuildMtpDraft). The verify /
  // propose runner loop and --speculative-config CLI parsing are I5d; this field
  // only enables the seam. Non-safetensors (GGUF) checkpoints lack `mtp.*`, so an
  // MTP config over a GGUF source is rejected.
  std::optional<vllm::SpeculativeConfig> speculative_config = std::nullopt;
};

// Owns the full V1 engine stack (config + weights + tokenizer + Scheduler +
// runner -> Executor -> EngineCore; Input/OutputProcessor -> LLMEngine) for a
// registered model. The concrete weights/forward are held behind LoadedModel;
// the Scheduler / Executor / EngineCore / processors are arch-agnostic (they
// touch the runner only through ModelRunnerBase).
// Members are declared in dependency order so the LLMEngine's by-reference
// collaborators stay valid for this object's lifetime. NON-COPYABLE /
// NON-MOVABLE (the internal references would dangle) — always heap-hold behind a
// unique_ptr and hand out engine() by reference.
class LoadedEngine {
 public:
  // Build the stack from already-loaded model pieces. This is the shared seam:
  // FromModelDir() loads config/tokenizer/weights from disk then calls this, and
  // tests construct it directly with synthetic in-memory weights (no disk). The
  // pieces are moved into members that outlive every collaborator that
  // references them.
  LoadedEngine(HfConfig config, Qwen3_5MoeWeights weights,
               tok::Tokenizer tokenizer, const EngineParams& params);

  // DENSE-arch overload (27B). Identical stack; the runner runs the dense
  // Qwen3_5DenseModel::Forward over the dense weights instead of the MoE forward.
  LoadedEngine(HfConfig config, Qwen3_5DenseWeights weights,
               tok::Tokenizer tokenizer, const EngineParams& params);

  LoadedEngine(const LoadedEngine&) = delete;
  LoadedEngine& operator=(const LoadedEngine&) = delete;
  LoadedEngine(LoadedEngine&&) = delete;
  LoadedEngine& operator=(LoadedEngine&&) = delete;

  // Load config.json + tokenizer.json + *.safetensors from `model_dir` and build
  // the stack. Throws std::runtime_error on any load failure (bad path, missing
  // shards, unparseable config).
  static std::unique_ptr<LoadedEngine> FromModelDir(const std::string& model_dir,
                                                    const EngineParams& params);

  // Resolve the per-step token budget (max_num_batched_tokens) for chunked
  // prefill. An explicit EngineParams override wins; otherwise a PER-ARCH
  // default (see the definition in model_loader.cpp for the measurements):
  //  * DENSE arch: 2048 flat — vLLM's own scheduler default
  //    (DEFAULT_MAX_NUM_BATCHED_TOKENS = 2048, vllm/config/scheduler.py:42 @
  //    e24d1b24). The dense prefill is expensive per token; a bigger budget
  //    lets one giant mixed step run several full prompts' prefill and starves
  //    every decode stream behind it.
  //  * MoE arch: the GB10-tuned concurrency-aware budget (8192 at
  //    max_num_seqs >= 32, else 4096) — the cheap A3B expert prefill wants the
  //    bigger chunk.
  // Invariants (SchedulerConfig.verify_max_model_len,
  // vllm/config/scheduler.py:87): result >= max_num_seqs; the tiny-model
  // (max_model_len * max_num_seqs) ceiling is preserved. Exposed for testing
  // the default policy without a disk load.
  static int ResolveMaxNumBatchedTokens(const EngineParams& params,
                                        int max_model_len, bool is_dense_arch);
  static bool ResolveEnablePrefixCaching(const EngineParams& params,
                                         const ModelInfo& model_info);

  vllm::v1::LLMEngine& engine() { return engine_; }
  // Lazily start W2's EngineCoreProc + output-handler threads. Once created,
  // online/server callers use this frontend rather than the synchronous
  // LLMEngine over the same scheduler/executor.
  vllm::v1::AsyncLLM& async_engine();
  const tok::Tokenizer& tokenizer() const { return tokenizer_; }
  const HfConfig& config() const { return config_; }
  int max_model_len() const { return max_model_len_; }
  bool prefix_caching_enabled() const { return prefix_caching_enabled_; }
  // ENG-SGLANG-BEHAVIOR-FLAG SW3: the jump-forward enable, resolved ONCE at
  // construction from EngineParams::enable_jump_forward + the
  // VT_ENABLE_JUMP_FORWARD env override (vllm::v1::JumpForwardEnabled). Exposed so
  // the enablement gate can assert the C-ABI/C++/flag toggle took effect.
  bool jump_forward_enabled() const { return jump_forward_enabled_; }
  const vllm::v1::GPUModelRunner& runner() const { return runner_; }

  // KV-EXTERNAL-CACHE (LMCache): the wired external KV connector, or null when
  // none was configured. Exposed so the output-invariance gate can read the
  // prefill-tokens-saved / chunks-stored counters. Non-owning.
  vllm::v1::kv_offload::KVConnector* kv_connector() const {
    return kv_connector_.get();
  }

  // The async-scheduling enable-flip, resolved ONCE at construction (W3
  // ENG-ASYNC-SCHED). `async_scheduling_enabled()` is
  // AsyncSchedulingEnabled(SchedulerConfig::ResolveAsyncScheduling(
  // runner_.runner_supports_async())) — i.e. vLLM's default-ON-when-compatible
  // resolution (vllm/config/vllm.py:990-1038) gated on the MRV2 runner
  // advertising the async device path (VT_ASYNC_RUNNER), with the house
  // VT_ASYNC_SCHED=0 rollback applied. When ON, scheduler() is an AsyncScheduler
  // and max_concurrent_batches() is 2 (depth-2 batch queue, step_with_batch_queue);
  // OFF keeps the byte-identical synchronous Scheduler + depth-1. Since the
  // 2026-07-17 flip the production default (no env) resolves ON (VT_ASYNC_RUNNER
  // default ON → AsyncScheduler + mcb=2), mirroring vLLM; VT_ASYNC_RUNNER=0 or
  // VT_ASYNC_SCHED=0 roll it back in the same binary.
  bool async_scheduling_enabled() const { return async_scheduling_enabled_; }
  int max_concurrent_batches() const { return max_concurrent_batches_; }
  // The engine's scheduler (Scheduler or, under async scheduling, AsyncScheduler).
  // Exposed for the construction-resolution gate; production drives it via
  // engine()/async_engine().
  const vllm::v1::Scheduler& scheduler() const { return *scheduler_; }

  // ResolveAsyncEnabled: the construction-time resolution, factored out so the
  // CPU construction-matrix test can assert it directly over the
  // runner_supports_async x VT_ASYNC_SCHED matrix without a disk load. Applies
  // SchedulerConfig::ResolveAsyncScheduling then the VT_ASYNC_SCHED rollback env.
  static bool ResolveAsyncEnabled(const vllm::SchedulerConfig& scheduler_config,
                                  bool runner_supports_async);

 private:
  // Type-erased constructor used by FromModelDir and the concrete-weight
  // compatibility overloads above.
  LoadedEngine(HfConfig config, std::unique_ptr<LoadedModel> model,
               tok::Tokenizer tokenizer, const EngineParams& params,
               vt::Queue* preselected_queue = nullptr,
               std::unique_ptr<DflashDraft> dflash_draft = nullptr);

  static vllm::SchedulerConfig MakeSchedulerConfig(
      int max_model_len, int max_num_seqs, int max_num_batched_tokens,
      vllm::SchedulerPolicy policy = vllm::SchedulerPolicy::kFCFS);
  // Construct the concrete scheduler for the resolved mode: an AsyncScheduler
  // (async scheduling ON) or the synchronous Scheduler. Mirrors upstream
  // get_scheduler_cls (scheduler.py:180-189) selecting AsyncScheduler when
  // async_scheduling. Both take the same ctor arguments.
  static std::unique_ptr<vllm::v1::Scheduler> MakeScheduler(
      bool async_enabled, vllm::SchedulerConfig scheduler_config,
      vllm::v1::KVCacheConfig kv_cache_config, int block_size,
      bool enable_caching,
      vllm::v1::StructuredOutputManager* structured_output_manager,
      std::optional<vllm::SpeculativeConfig> speculative_config = std::nullopt);
  // SPEC-MTP I5d: finalize the entrypoint's SpeculativeConfig against the loaded
  // checkpoint. params.speculative_config carries the CLI method + optional user
  // k; this re-runs SpeculativeConfig::ResolveMtp with the checkpoint's
  // mtp_num_hidden_layers (from config.raw text_config, default 1) so n_predict
  // and the resolved k are correct. Returns nullopt when no speculation is
  // configured (the byte-identical production path). Non-Qwen3.5 or non-mtp
  // methods that reached here throw.
  static std::optional<vllm::SpeculativeConfig> ResolveSpecConfig(
      const EngineParams& params, const HfConfig& config);
  // SPEC-MTP I5d: build the KV-cache spec, widened for speculation when a spec
  // config is set (the extra GDN k+1 state slots + widened conv row + the
  // `fa_draft` full-attn group, MakeQwen3_5KVCacheSpec num_spec>0). With no spec
  // config it is exactly ModelRegistry::MakeKVCache (byte-identical).
  static vllm::v1::KVCacheConfig MakeKVCacheMaybeSpec(
      const LoadedModel& model, const HfConfig& config, int block_size,
      int num_blocks, const std::optional<vllm::SpeculativeConfig>& spec);
  // Ensure NONE_HASH is initialized before the scheduler/hasher are built
  // (upstream global init). Idempotent; runs as the first member initializer.
  static bool EnsureNoneHash();
  // Mirrors kernel_warmup.py::flashinfer_autotune: before any async/server
  // frontend starts, one maximum-token synthetic request runs under the NVFP4
  // all-bucket autotune scope. CUDA dense W4A4 only; CPU/other models are no-op.
  void WarmupKernels();

  bool hash_ready_;  // declared first: forces EnsureNoneHash() ahead of the rest.
  HfConfig config_;
  // SPEC-MTP I5d: the finalized speculative config (method/k/n_predict), or
  // nullopt on the production default path. Declared before model_/kv_cfg_/runner_
  // because the KV-cache widening, the draft build, the scheduler lookahead, and
  // the engine-core draft pull all read it. nullopt ⇒ every spec path is inert
  // and the engine is byte-identical to the pre-spec engine.
  std::optional<vllm::SpeculativeConfig> resolved_spec_config_;
  // SPEC-DFLASH D5: the owned DFlash draft (null unless method=="dflash").
  // Declared before runner_ so the borrow the runner holds stays live for the
  // runner's whole lifetime. Set in the ctor body via runner_.set_dflash_draft.
  std::unique_ptr<DflashDraft> dflash_draft_;
  // Concrete weights and model-specific runtime state behind the central
  // registry contract. Declared before runner_ so its borrow remains live.
  std::unique_ptr<LoadedModel> model_;
  // KV-EXTERNAL-CACHE (LMCache): the owned external KV connector (null unless
  // EngineParams::kv_transfer_config selects one). Declared BEFORE runner_ /
  // scheduler_ so it outlives the non-owning pointers they hold to it. Built in
  // the ctor body once runner_ geometry is known; see model_loader.cpp.
  std::unique_ptr<vllm::v1::kv_offload::KVConnector> kv_connector_;
  tok::Tokenizer tokenizer_;
  int max_model_len_;
  int max_num_batched_tokens_;
  bool prefix_caching_enabled_;
  // ENG-SGLANG-BEHAVIOR-FLAG SW3: jump-forward enable, resolved once from
  // EngineParams::enable_jump_forward + the VT_ENABLE_JUMP_FORWARD env override.
  // Depends only on params + env (no member deps), so its init order is free.
  bool jump_forward_enabled_;
  vllm::v1::KVCacheConfig kv_cfg_;
  // runner_ is declared BEFORE the scheduler (W3): the async-scheduling flip is
  // resolved from runner_.runner_supports_async(), so the runner must be fully
  // constructed before async_scheduling_enabled_ / scheduler_ are initialized.
  vllm::v1::GPUModelRunner runner_;
  // Resolved once, in dependency order after runner_ (see the accessors above).
  bool async_scheduling_enabled_;
  int max_concurrent_batches_;
  // The engine's StructuredOutputManager (upstream EngineCore constructs one,
  // core.py:134), wired with the NATIVE backend factory over tokenizer_ (which
  // outlives it — declared earlier) so response_format / C-ABI structured
  // constraints actually gate decoding. Declared before scheduler_ /
  // engine_core_, which hold a pointer to it.
  vllm::v1::StructuredOutputManager structured_output_manager_;
  // Polymorphic scheduler: a plain Scheduler by default, an AsyncScheduler when
  // async_scheduling_enabled_. Heap-held so the concrete class is chosen at
  // construction; engine_core_ / async_engine_ share this one instance by
  // reference (*scheduler_).
  std::unique_ptr<vllm::v1::Scheduler> scheduler_;
  vllm::v1::Executor executor_;
  vllm::v1::EngineCore engine_core_;
  vllm::v1::InputProcessor input_processor_;
  vllm::v1::OutputProcessor output_processor_;
  vllm::v1::BlockHasher block_hasher_;
  vllm::v1::LLMEngine engine_;
  std::mutex async_engine_mutex_;
  std::unique_ptr<vllm::v1::AsyncLLM> async_engine_;
};

}  // namespace vllm::entrypoints

#endif  // VLLM_ENTRYPOINTS_MODEL_LOADER_H_