vllm-project / vllm-project/vllm-neuron

Add Qwen3.5 support (hybrid gated-DeltaNet + attention decoder): dense 2B/27B, then sparse 35B-A3B

Open
#53 2 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

bug
Dominant language
Python
Stars
51
Forks
35
PR merge metrics
No merged PRs in 30d

Description

Summary

I have Qwen3.5 running on this plugin at release-0.24.0.1.1.0 — the dense
checkpoints (2B, 27B) and, separately, the sparse 35B-A3B — and would like to
contribute them upstream. Opening this first per CONTRIBUTING ("open an issue to
discuss any significant work"), because the dense half touches shared code and
I would rather find out now if the approach is wrong.

What Qwen3.5 needs that the plugin does not have yet

Qwen3.5's decoder is hybrid: most layers are gated DeltaNet, a linear
recurrence carrying fixed-size state per sequence, and the rest are ordinary
full attention (18 + 6 on 2B, 48 + 16 on 27B). vLLM puts the two kinds of layer
in separate KV cache groups, which this plugin has not had to support before.
Four things follow:

  1. A model needs a way to declare state-carrying layers. KVSpec/LayerSpec
    describe paged attention layers only; num_kv_heads/head_size are
    meaningless for a recurrent layer, where what matters is the concrete state
    tensor shapes.

  2. Page sizes have to be reconciled across the two groups.
    unify_kv_cache_spec_page_size requires every group's page size to divide the
    largest. A DeltaNet state page is 271360 bytes per rank at TP=4, which factors
    as 1024 * 5 * 53, so no sane attention block size divides it and startup
    fails with "Cannot unify by adjusting block_size". vLLM solves this in
    Platform._align_hybrid_block_size — grow the attention block size, then pad
    the state page to match — but it gates that on _find_non_ssm_backend
    finding one of its own attention backends, and NeuronPlatform registers
    none, so it never runs here.

  3. Neuron device tensors reject vLLM's state layout. gpu_model_runner
    strides each state tensor out of the layer's raw buffer with the block dim
    striding by a whole page. On Neuron that raises Detected non-contiguous slicing for requested Device Tensor for any view whose stride is not the
    natural one, so the layout has to be chosen differently — and the choice
    interacts with whether the raw buffer is shared across groups.

  4. --modular-flow-mac-threshold=10 breaks codegen on a pure-torch graph.
    The runner passes it unconditionally. On this model's decode graph neuronx-cc
    fails with NCC_IBTN006 (a pftranspose whose copy fails backend
    verification). Confirmed narrowly by recompiling the cached graph.hlo by
    hand: fails with the flag, compiles without it. The threshold exists because
    NKI kernels do not report MAC counts, and the dense path calls no kernels.

What I propose

A PR against release-0.24.0.1.1.0 (not main — the code is written against
0.24.0.1.1.0's neuron_model_runner/platform surface, and recent PRs target
that branch) adding:

  • vllm_neuron/model/qwen3_5/ — the model, and Qwen3_5ForConditionalGeneration
    in the registry.
  • RecurrentLayerSpec + KVSpec.recurrent_layers in model/kv_cache.py.
  • MambaSpec construction and both cache layouts in the model runner.
  • _align_hybrid_page_sizes in platform.py, which calls vLLM's own helper with
    a stub backend supplying the one thing it reads, rather than reimplementing the
    arithmetic — the same helper sizes the state pages the planner allocates, and a
    second copy of that arithmetic is a silent aliasing bug waiting to happen.
  • NeuronConfig.hlo2tensorizer_options, so a model can opt out of the threshold.
  • examples/vllm_neuron/models/qwen3_5/ with a run script, a latency benchmark
    and an HF-comparison accuracy check.

Measured on a trn2.3xlarge at TP=4, bf16, greedy, 1024 in / 128 out: 2B does
TTFT 145 ms / TPOT 3.88 ms at batch 1 and 465 output tok/s aggregate at batch 8.
27B runs with gpu_memory_utilization=0.65 and -O3; the reason for each is
documented in the example README, along with the known limitations (prefix
caching must stay off; one batch bucket per process; decode gathers the full
block-table width rather than only the occupied pages).

The sparse 35B-A3B support is a second, independent PR — a Qwen3_5SparseMoeBlock
plus a second registry entry — so that the dense change can be reviewed on its
own.

Attribution

The Qwen3.5 implementation derives from
qingzwang/vllm-neuron branch
model/Qwen3.5-2B (Apache-2.0), which is where the 2B port originates;
nki_deltanet.py, nki_deltanet_fused.py, vl.py and flags.py are that
fork's files unchanged, and the bulk of model.py, deltanet.py and config.py
is its work. My contribution is the forward-port to 0.24.0.1.1.0, the 27B
checkpoint, and the sparse variant. Credited in the commit trailer as well.

Questions

  1. Is release-0.24.0.1.1.0 the right base, or would you prefer main?
  2. Is extending KVSpec with a separate recurrent_layers list the shape you
    want for hybrid models, or would you rather LayerSpec grew a kind
    discriminator?
  3. There is no test directory on this branch. Where should model tests go?

Happy to split, rebase or restructure — please say which.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start on release-0.24.0.1.1.0 by reviewing model/kv_cache.py, platform.py, neuron_model_runner, and NeuronConfig.hlo2tensorizer_options. Compare the proposed hybrid cache and page-size handling with vLLM's _align_hybrid_block_size, then inspect the listed qwen3_5 entry points and examples. Done means dense 2B and 27B support is registered, runnable, and covered by the proposed accuracy and benchmark checks.

Written by the indexing model from the issue text.

Assessment

Tech stack
aws, python, pytorch
Domain
ai-infra-agents, infrastructure
Issue type
Feature
Difficulty
5/5
Estimated time
Over a week
Activity status
Active
Clarity
Mostly clear
Newbie friendliness
35/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.