Add Qwen3.5-2B (support vision and text input)

[qingzwang]

Description

Add Qwen3.5-2B contrib model — NeuronX Distributed Inference port of Qwen/Qwen3.5-2B, a 2B-parameter vision-language decoder with a hybrid attention stack: 24 transformer layers arranged as [3 gated DeltaNet (linear attention) + 1 full GQA attention] × 6.

This is the first hybrid linear-attention + softmax-attention model in the contrib folder. It supports both text-only and vision-language inference end-to-end on trn2.48xlarge.

Key implementation highlights:

• Hybrid DeltaNet + GQA architecture — 18 gated DeltaNet layers (recurrent, O(1) state) + 6 full GQA layers with per-head query gating and partial RoPE (2D + interleaved mRoPE, head_dim=256, partial_rotary_factor=0.25)
• Custom NKI kernels for DeltaNet: fused chunked forward (CTE path), recurrent step, chunked step, legacy-direct variant (required for VL), and optional QK-norm+RoPE kernel
• Vision-language support: ViT encoder (24 layers, patch_size=16) compiled for Neuron with TP-sharding (2× TTFT speedup vs CPU baseline) + Neuron text decoder. Includes CTE bucketing ([1024, 4096]), 2×2 spatial tiling for 2048×2048 images, get_rope_index 3-D mRoPE, and index_put_ scatter at image-token positions. Also supports CPU-only VE fallback.
• No NxDI library modifications required — runs on stock DLAMI venv
• Reuses DeltaNet + GQA modeling code from PR #173 (27B sibling), adapted for the 2B variant
• Validated at both TP=8 (primary) and TP=4

Model Information

Model Name: Qwen3.5-2B

Model Architecture: Hybrid decoder — [3 gated DeltaNet (linear attention) + 1 GQA (softmax attention)] × 6, with ViT vision encoder

Purpose: Text generation and vision-language (image-to-text) inference

Checklist

Required Components

• Accuracy Test (test/integration/test_model.py)

  • Integration test with logit validation and greedy generation comparison
  • 66% exact-token greedy match vs. HuggingFace CPU bf16 reference (3/5 prompts fully match; divergence is expected bf16 accumulation difference)
  • Additional tests: run_text_smoke.py, run_benchmark.py, run_accuracy_check.py, run_hf_reference.py, run_vl_smoke.py, run_vl_benchmark.py, compile_vision_encoder.py, compile_vision_encoder_tp.py

• README.md with the following sections:

  • Usage Example: Full CLI command for text-only smoke test with expected output
  • Compatibility Matrix: Detailed component version table (neuronx-cc, nki, nxd, nxdi, torch, Python)
  • Example Checkpoints: Link to Qwen/Qwen3.5-2B (~4.5 GB bf16)
  • Testing Instructions: Commands for text smoke, benchmark, accuracy check, and VL smoke

• Source Code (src/)

  • modeling_qwen35.py — text decoder (DeltaNet + GQA + MRoPE, ~324 KB)
  • modeling_qwen35_vision.py — ViT encoder wrapper (~40 KB)
  • modeling_qwen35_vl.py — VL orchestrator (~28 KB)
  • hybrid_apc.py — Automatic Prefix Cache (~70 KB, unused by default)
  • nki_kernels/ — 5 NKI kernel files for DeltaNet

Optional Components

• Unit Tests (test/unit/ directory exists)

Folder Structure

/contrib/models/Qwen3.5-2B/
  README.md
  /src
    __init__.py
    modeling_qwen35.py
    modeling_qwen35_vision.py
    modeling_qwen35_vl.py
    hybrid_apc.py
    /nki_kernels
      __init__.py
      nki_deltanet.py
      nki_deltanet_chunked.py
      nki_deltanet_fused.py
      nki_deltanet_fused_legacy.py
      qwen_qk_norm_rope.py
  /test
    __init__.py
    /unit
    /integration
      __init__.py
      test_model.py
      run_text_smoke.py
      run_benchmark.py
      run_hf_reference.py
      run_accuracy_check.py
      run_vl_smoke.py
      run_vl_benchmark.py
      compile_vision_encoder.py
      compile_vision_encoder_tp.py

Testing

How did you test this change?

Tested on trn2.48xlarge with:

• Text-only: compile + generate at TP=8 (seq_len=512) and TP=4 (seq_len=1024)
• Vision-language: CPU vision encoder → Neuron text decoder on real images
• Accuracy cross-check against HuggingFace transformers==5.13.0 CPU bf16 greedy
• Benchmark sweep across prompt lengths (16, 64, 128, 256, 512 tokens)

Test Results:

Text-only performance (TP=8, bf16, seq_len=512, batch=1):

Prompt tokens	TTFT (ms)	TPOT (ms)	Throughput (tok/s)
16	17.6	4.00	250.4
64	17.6	3.99	250.2
256	17.5	4.00	250.7

TTFT is essentially flat because DeltaNet is O(1)-state and only 6/24 layers are full-attention.

TP=4 comparison (seq_len=1024):

Prompt tokens	TTFT (ms)	TPOT (ms)	Throughput (tok/s)
16–512	~42	~4.75	~210

Vision-language (TP=8, Neuron VE with CTE bucketing):

Image	Neuron VE TTFT (ms)	vs CPU VE	TPOT (ms)
512×512	86	2.0× faster	4.2
1024×1024	526	2.1× faster	4.1
2048×2048	2,090 (2×2 tiled)	2.4× faster	4.2

Optional TP-sharded vision encoder (compile_vision_encoder_tp.py) further parallelizes the ViT across multiple NeuronCores via parallel_model_trace(tp_degree=N).

Accuracy vs. HF reference: 53/80 = 66% exact-token greedy match (3/5 prompts fully match). Qualitatively coherent on all prompts.

Compatibility

Tested with:

• Neuron SDK: neuronx-cc 2.26.6360, NKI 0.5.0
• Instance Type(s): trn2.48xlarge (TP=8 primary, TP=4 also validated)
• NxDI Version: 0.10.18399
• NxD Version: 0.19.28492
• PyTorch Version: 2.9.1 (torch-neuronx 2.9.0.2)
• Python Version: 3.12

Additional Information

• First hybrid linear-attention model in the contrib folder — introduces DeltaNet (gated linear attention with delta rule) alongside standard GQA softmax attention.
• Custom NKI kernels implement the DeltaNet recurrence efficiently on NeuronCores. The fused chunked kernel is the default CTE path; the legacy-direct variant is needed for VL (variable-length vision token sequences).
• No NxDI library modifications — runs on the stock /opt/aws_neuronx_venv_pytorch_2_9_nxd_inference/ DLAMI venv.
• Vision encoder runs on Neuron — compiled with CTE bucketing ([1024, 4096]), 2× TTFT speedup vs CPU baseline. Optional TP-sharded mode (compile_vision_encoder_tp.py) distributes QKV/MLP across cores. 2048×2048 images use transparent 2×2 spatial tiling through the 4096-bucket kernel.
• CPU VE fallback available — run_vl_smoke.py can fall back to CPU vision (cosine similarity 0.99 vs. HF reference) for debugging.
• Relationship to PR Add Qwen3.6-27B hybrid DeltaNet/GQA contrib model and vLLM serving path #173 — this 2B variant reuses the DeltaNet + GQA modeling code contributed for the 27B sibling (Qwen3.5-27B).
• HF reference requires transformers ≥ 5.13 — the qwen3_5 architecture identifier was introduced in transformers 5.x (the Neuron runtime only needs 4.57+).
• VL uses legacy-direct DeltaNet kernel — the fused-multihead NKI kernel is numerically unstable on real vision embeddings; the legacy-direct kernel is stable and is the default for VL inference.

Related Issues

Related to PR #173 (Qwen3.5-27B DeltaNet + GQA contribution).

vLLM Integration

• This model/feature is intended for use with vLLM
• Documentation includes vLLM registration instructions

vLLM integration is future work.

---

By submitting this PR, I confirm that:

• I have read and followed the contributing guidelines
• This is a community contribution and may have limited testing compared to officially-supported models
• The code follows best practices and is well-documented
• All required components listed above are included