qwen3_5_dense.cpp

// Qwen3.6 DENSE (27B) architecture registry TU. Self-registers
// "Qwen3_5ForConditionalGeneration" via REGISTER_VLLM_MODEL and owns the dense
// arch-specific entry points (LoadedModel subclass + load/prepare/forward
// wrappers + factory + synthetic Make/Borrow adapters). The heavy dense forward
// machinery (Qwen3_5DenseModel::/Qwen3_5DenseDecodeGraph::) lives in qwen3_5.cpp
// over the shared DevicePool/matmul/GDN helpers; this TU only wires it into the
// registry. Extracted verbatim (behavior-preserving) from the former
// model_registry.cpp monolith.
#include "vllm/v1/worker/gpu/cudagraph_dispatch.h"
#include "vllm/model_executor/models/model_registry.h"

#include <cstdlib>
#include <memory>
#include <optional>
#include <stdexcept>
#include <utility>

#include "vllm/model_executor/models/qwen3_5.h"         // ForwardLogits
#include "vllm/model_executor/models/qwen3_5_common.h"  // kQwen3_5Info, helpers
#include "vllm/model_executor/models/qwen3_5_dense.h"
#include "vllm/model_executor/models/qwen3_5_gguf_weights.h"
#include "qwen3_5_internal.h"  // W4 DeviceTokenIdsScope
#include "vllm/model_executor/models/qwen3_5_mtp.h"  // SPEC-MTP I5d-pre draft
#include "vllm/platforms/interface.h"  // GetPlatform(device.type) memory-model seam

namespace vllm {
namespace {

bool DenseDecodeGraphEnabled() {
  const char* value = std::getenv("VLLM_CPP_DENSE_DECODE_GRAPH");
  return value == nullptr || value[0] != '0';
}

class Qwen3_5DenseLoadedModel final : public LoadedModel {
 public:
  Qwen3_5DenseLoadedModel(const ModelRegistration& registration,
                          Qwen3_5DenseWeights weights)
      : LoadedModel(registration),
        owned_weights_(std::move(weights)),
        weights_(&*owned_weights_) {}
  Qwen3_5DenseLoadedModel(const ModelRegistration& registration,
                          const Qwen3_5DenseWeights& weights,
                          BorrowedWeightsTag)
      : LoadedModel(registration), weights_(&weights) {}

  const Qwen3_5DenseWeights& weights() const { return *weights_; }
  bool uses_nvfp4_w4a4() const override {
    return !weights_->layers.empty() &&
           weights_->layers.front().mlp.gate_proj_fp4.IsTrueW4A4();
  }
  std::unique_ptr<Qwen3_5DenseDecodeGraph>& decode_graph() {
    return decode_graph_;
  }

  // SPEC-MTP I5d-pre: retain the loaded `mtp.*` draft weights + build the draft
  // sharing this target's embed_tokens/lm_head. Inert unless FromModelDir
  // attached weights (i.e. unless a SpeculativeConfig is configured).
  bool supports_mtp_draft() const override { return true; }
  // SPEC-DFLASH / SPEC-DSPARK: this forward routes to ForwardDeviceMultiTap.
  bool supports_aux_multi_tap() const override { return true; }
  void AttachMtpDraftWeights(Qwen3_5MTPWeights weights) override {
    mtp_draft_weights_ = std::move(weights);
  }
  std::unique_ptr<Qwen3_5MTPModel> BuildMtpDraft(
      const HfConfig& config) const override {
    if (!mtp_draft_weights_.has_value()) return nullptr;
    return std::make_unique<Qwen3_5MTPModel>(*mtp_draft_weights_, *weights_,
                                             config);
  }

 private:
  std::optional<Qwen3_5DenseWeights> owned_weights_;
  const Qwen3_5DenseWeights* weights_ = nullptr;
  std::unique_ptr<Qwen3_5DenseDecodeGraph> decode_graph_;
  // Retained draft weights (SPEC-MTP I5d-pre); empty on the production default.
  std::optional<Qwen3_5MTPWeights> mtp_draft_weights_;
};

std::unique_ptr<LoadedModel> LoadQwen3_5DenseModel(
    const ModelRegistration& registration, const HfConfig& config,
    const ModelSource& source) {
  if (source.kind == ModelSource::Kind::kGguf) {
    // Dense-arch (`qwen35`, e.g. Qwen3.5-2B) GGUF: same GDN / full-attention
    // block loaders as the MoE path with a dense SwiGLU MLP per layer.
    if (source.gguf == nullptr) {
      throw std::runtime_error("GGUF model source is empty");
    }
    return std::make_unique<Qwen3_5DenseLoadedModel>(
        registration, LoadQwen3_5DenseFromGguf(*source.gguf, config));
  }
  if (source.kind != ModelSource::Kind::kSafetensors) {
    throw std::runtime_error(
        "Model architecture Qwen3_5ForConditionalGeneration does not support "
        "this weight source");
  }
  if (source.safetensors == nullptr) {
    throw std::runtime_error("safetensors model source is empty");
  }
  return std::make_unique<Qwen3_5DenseLoadedModel>(
      registration,
      LoadQwen3_5Dense(*source.safetensors, config, source.load_queue));
}

void PrepareQwen3_5Dense(LoadedModel& model, const HfConfig& config,
                         vt::Queue& queue) {
  // PERF-27B-LMHEAD-FP4 (issue #213): build the packed lm_head's resident HERE —
  // on CUDA before the runner captures a decode graph, elsewhere before the first
  // forward pays the dequant. Inert on every BF16/FP8/GGUF/tied checkpoint.
  auto& qwen = static_cast<Qwen3_5DenseLoadedModel&>(model);
  Qwen3_5DenseModel::PrepareLmHeadResident(qwen.weights(), queue);
  // PERF-27B-GDN-FP8-QKVZ: build the merged FP8 GDN [qkv;z] operand here, at
  // model prepare — before the first forward, so it can never allocate or copy
  // inside a CUDA-graph capture. No-op on CPU, on a non-FP8 owner, and when the
  // merge is rolled back (VT_GDN_MERGED_QKVZ_FP8=0).
  Qwen3_5DenseModel::PrepareGdnFp8Resident(qwen.weights(), config, queue);
}

ForwardLogits ForwardQwen3_5Dense(LoadedModel& model,
                                  const ModelForwardInput& input) {
  auto& qwen = static_cast<Qwen3_5DenseLoadedModel&>(model);
  const Qwen3_5DenseWeights& weights = qwen.weights();

  // ENG-ASYNC-SCHED W4: publish the async runner's device-resident input ids for
  // the duration of THIS forward, so the embed at the top of every route below
  // (eager, gathered, tap, multi-tap, decode-graph replay) reads them instead of
  // uploading the host vector, which is stale for decode rows on that path.
  // Null on every other path, and RAII-scoped so it cannot outlive the call.
  const detail::DeviceTokenIdsScope device_ids_scope(
      input.device_token_ids, static_cast<int64_t>(input.token_ids.size()));

  // SPEC-MTP I5d-pre hidden-state tap. When the spec verify forward requests the
  // drafter's [T,H] post-final-norm hidden (I5d), route to the EXISTING
  // ForwardDeviceTap: byte-identical logits to ForwardDevice, plus the hidden
  // moved into *input.hidden_tap. Null (every spec-off run) falls through to the
  // unchanged path below, so the forward is byte-identical when spec is off.
  if (input.hidden_tap != nullptr) {
    return Qwen3_5DenseModel::ForwardDeviceTap(
        input.token_ids, input.positions, input.attn_meta, input.gdn_meta,
        input.attn_kv, input.gdn_state, weights, input.config, input.queue,
        input.hidden_tap, input.logits_indices);
  }

  // SPEC-DFLASH D1 (DF-AUX-TAPS): non-null routes to ForwardDeviceMultiTap
  // (byte-identical logits + the [T,H×taps] aux capture); null is byte-identical to
  // the path below. Mutually exclusive with hidden_tap.

  // vLLM's CUDA-graph selection is independent of weight quantization. The
  // driver below already captures the shared dense forward, so restricting it
  // to the 27B FP4 checkpoint left ordinary BF16 Qwen3.5 decode eager.
  const bool graph_cuda =
      platforms::GetPlatform(input.queue.device.type).support_static_graph_mode();
  constexpr int kMaxDecodeGraphBatch = 64;

  // SPEC-DSPARK W8 (#442): mirror vLLM's UNIFORM-decode predicate instead of
  // "query_len == 1". Upstream's captured decode length is
  // `1 + num_speculative_tokens` (cudagraph_dispatcher.py:37), so its T=1+k
  // speculative VERIFY is graph-captured; ours fell to the eager path EVERY step,
  // which the paired 35B measurement charged at ~4.8 ms/step (0.870x where
  // acceptance is high). With num_speculative_tokens == 0 this is exactly
  // `pure_decode`, so the non-spec path is unchanged.
  // DEFAULT ON. The staging defect that forced this off is fixed: the captured
  // graph reads gdn_spec_* / num_accepted, which StageSpecStepInputs now refills
  // per step, and the per-request arrays are sized by the REQUEST count rather
  // than the token count. Gated green with capture ON and OFF across the MTP,
  // DFlash, 35B and concurrent e2e suites, and measured +8.5% / +4.7% on the 35B
  // cells (0.870x -> 0.986x of the pinned graphed oracle on the high-acceptance
  // one). VT_SPEC_DECODE_GRAPH=0 restores the eager verify for an A/B.
  static const bool spec_graph = [] {
    const char* v = std::getenv("VT_SPEC_DECODE_GRAPH");
    return v == nullptr || (v[0] != '\0' && v[0] != '0');
  }();
  const bool uniform_decode =
      input.pure_decode ||
      (spec_graph && input.num_speculative_tokens > 0 && input.gdn_meta.num_prefill_tokens == 0 &&
       v1::IsUniformDecodeBatch(input.num_reqs, input.attn_meta.num_actual_tokens,
                                input.attn_meta.max_query_len,
                                input.num_speculative_tokens) &&
       // The captured region emits EVERY row, so only a no-op gather may take it.
       // A spec verify needs all 1+k rows anyway, so this holds there.
       (input.logits_indices.empty() ||
        static_cast<int64_t>(input.logits_indices.size()) ==
            input.attn_meta.num_actual_tokens));
  if (DenseDecodeGraphEnabled() && uniform_decode && graph_cuda &&
      input.num_reqs <= kMaxDecodeGraphBatch) {
    if (!qwen.decode_graph()) {
      qwen.decode_graph() = std::make_unique<Qwen3_5DenseDecodeGraph>(
          weights, input.config, input.queue, input.gdn_state_slots);
    }
    return qwen.decode_graph()->Step(
        input.token_ids, input.positions, input.attn_meta, input.gdn_meta,
        input.attn_kv, input.gdn_state, input.aux_tap);
  }

  // SPEC-DSPARK W8 (#442): the aux multi-tap forward is the DFlash/DSpark
  // VERIFY path, and it used to return BEFORE the decode-graph gate, which is
  // why the T=1+k verify never captured. The graph branch above now serves it
  // (writing the taps into the slot's persistent buffer); this remains the
  // eager fallback for every batch the graph declines.
  if (input.aux_tap != nullptr) {
    return Qwen3_5DenseModel::ForwardDeviceMultiTap(
        input.token_ids, input.positions, input.attn_meta, input.gdn_meta,
        input.attn_kv, input.gdn_state, weights, input.config, input.queue,
        input.aux_tap, input.logits_indices);
  }
  if (input.gather_logits) {
    return Qwen3_5DenseModel::ForwardDevice(
        input.token_ids, input.positions, input.attn_meta, input.gdn_meta,
        input.attn_kv, input.gdn_state, weights, input.config, input.queue,
        input.logits_indices);
  }
  return HostLogits(
      Qwen3_5DenseModel::Forward(
          input.token_ids, input.positions, input.attn_meta, input.gdn_meta,
          input.attn_kv, input.gdn_state, weights, input.config, input.queue,
          input.logits_indices),
      input.config.vocab_size);
}

const ModelFactory kQwen3_5DenseFactory{
    .parse_config = &ParseQwen3_5Config,
    .load_weights = &LoadQwen3_5DenseModel,
    .prepare = &PrepareQwen3_5Dense,
    .forward = &ForwardQwen3_5Dense,
    .make_kv_cache = &MakeQwen3_5KVCache,
    .is_dense_model = true,
};

}  // namespace

std::unique_ptr<LoadedModel> MakeQwen3_5DenseLoadedModel(
    Qwen3_5DenseWeights weights) {
  return std::make_unique<Qwen3_5DenseLoadedModel>(
      RegistrationFor("Qwen3_5ForConditionalGeneration"), std::move(weights));
}

std::unique_ptr<LoadedModel> BorrowQwen3_5DenseLoadedModel(
    const Qwen3_5DenseWeights& weights) {
  return std::make_unique<Qwen3_5DenseLoadedModel>(
      RegistrationFor("Qwen3_5ForConditionalGeneration"), weights,
      BorrowedWeightsTag{});
}

REGISTER_VLLM_MODEL(qwen3_5_dense, "Qwen3_5ForConditionalGeneration",
                    kQwen3_5DenseFactory, kQwen3_5Info)

}  // namespace vllm