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

[qingzwang]

Description

Add Qwen3.5-35B-A3B contrib model — NeuronX Distributed Inference port of Qwen/Qwen3.5-35B-A3B, the MoE flagship of the Qwen3.5 family. 35B total parameters, ~3B activated per token through top-8 routing over 256 experts + 1 sigmoid-gated shared expert, with the same hybrid attention stack as the dense siblings: [3 gated DeltaNet + 1 full GQA] × 10 = 40 layers.

This is the first DeltaNet + MoE integration on Neuron. It reuses the DeltaNet + attention path from PR #173 (Qwen3.5-27B dense) and integrates NxDI's initialize_moe_module (from moe_v2) via a new Qwen35MoEBlock. Both text-only and vision-language inference are validated end-to-end on trn2.48xlarge.

Key implementation highlights:

• Qwen35MoEBlock wraps NxDI's initialize_moe_module for 256 routed experts + a per-token sigmoid-gated shared expert (NxDI's built-in SharedExperts lacks per-token gating)
• Weight converter handles HF stacked 3D expert weights → NxDI's ExpertMLPsV2 transposed layout
• Three MoE blockwise-matmul backends evaluated: shard_on_intermediate_dynamic_while (default, 13% faster at short prompts), use_torch_block_wise, and monkey-patched _call_shard_hidden_kernel
• DeltaNet NKI kernels — unchanged from dense siblings (fused chunked forward CTE)
• Neuron VL — Neuron-compiled vision encoder (buckets 1024/4096) + text model with CTE bucketing, 2×2 spatial tiling for 2048×2048 images
• No NxDI library modifications required — runs on stock DLAMI venv

Model Information

Model Name: Qwen3.5-35B-A3B

Model Architecture: Hybrid DeltaNet + GQA decoder with Mixture-of-Experts FFN (256 routed + 1 shared, top-8 sigmoid routing, ~3B active/token), 40 layers, model_type=qwen3_5_moe_text

Purpose: Text generation and vision-language inference

Checklist

Required Components

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

  • All 5 prompts produce coherent, factually correct Neuron output
  • HF greedy comparison deferred (67 GB model too large for CPU bf16 in reasonable time)

• README.md with the following sections:

  • Usage Example: Full CLI commands for text-only and VL
  • Compatibility Matrix: Detailed component version table
  • Example Checkpoints: Link to Qwen/Qwen3.5-35B-A3B (~67 GB bf16)
  • Testing Instructions: Commands for text smoke, benchmark, VL compile + benchmark

• Source Code (src/)

  • modeling_qwen35.py — DeltaNet + GQA + NEW Qwen35MoEBlock (~336 KB)
  • modeling_qwen35_vision.py — ViT encoder wrapper
  • modeling_qwen35_vl.py — VL orchestrator
  • hybrid_apc.py — Automatic Prefix Cache
  • nki_kernels/ — DeltaNet NKI kernels (unchanged from dense siblings)

Optional Components

• Unit Tests (test/unit/ directory)

Folder Structure

/contrib/models/Qwen3.5-35B-A3B/
  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

Testing

How did you test this change?

Tested on trn2.48xlarge (TP=8, bf16) with:

• Text-only: compile + generate (5 diverse prompts, all coherent + factually correct)
• Benchmark: TTFT/TPOT sweep across prompt lengths with 3 MoE backends
• VL: Neuron vision encoder (buckets 1024/4096) + text model, tested across 512×512, 1024×1024, 2048×2048 images
• Correctly identifies Pallas's cat (manul) in test images

Test Results:

Text-only performance (TP=8, bf16, seq_len=512, shard_on_intermediate default):

Prompt tokens	TTFT (ms)	TPOT (ms)	Throughput (tok/s)
16	480.0	~7.7	~130
64	575.5	~7.7	~130
256	667.1	~7.7	~130

TPOT of ~7.7 ms matches the ~3B activated parameter budget (sits between 9B at 6.9ms and 4B at 5.7ms from dense siblings).

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

Image	TTFT (ms)	TPOT (ms)	tok/s
512×512	1,037	8.2	122
1024×1024	2,004	8.1	124
2048×2048 (2×2 tiled)	5,490	8.1	123

Compatibility

Tested with:

• neuronx-cc: 2.26.6360.0
• NKI: 0.5.0
• Instance Type: trn2.48xlarge (TP=8)
• NxDI: 0.10.18399
• NxD: 0.19.28492
• PyTorch: 2.9.1 (torch-neuronx 2.9.0.2)
• Python: 3.12
• transformers: 4.57.6 (HF CPU reference needs ≥ 5.13)

Additional Information

• First DeltaNet + MoE model on Neuron — combines hybrid linear-attention architecture with sparse MoE routing on NxDI.
• MoE config: MoENeuronConfig (not NeuronConfig), moe_tp_degree=8, moe_ep_degree=1, router.dtype=float32, router.act_fn=softmax, normalize_top_k_affinities=True.
• Shipped SDK MoE kernel gap: The preferred _call_shard_hidden_kernel LNC=2 path is not present in shipped neuronx-cc 2.26.6360; current default (shard_on_intermediate_dynamic_while) is 13% faster than torch fallback at short prompts but scales linearly with prompt length. A future SDK drop should close this.
• Expert Parallelism (EP=1): moe_ep_degree > 1 works for prefill but NxDI's ExpertMLPsV2 raises NotImplementedError for selective-loading decode with EP when top_k/num_experts < threshold.
• Per-token sigmoid-gated shared expert: Qwen3.5-MoE specific — NxDI's built-in SharedExperts only sums without per-token gating, so we implement it directly in Qwen35MoEBlock.
• HF greedy comparison deferred — 67 GB model too large for CPU bf16 in reasonable time; all Neuron outputs are qualitatively correct.
• VL uses legacy-direct DeltaNet kernel — same as dense siblings for numerical stability.

Related Issues

Related to PR #173 (Qwen3.5-27B DeltaNet + GQA dense 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