danym/qwen3.5-9b-demoted-mi300x-262k

Qwen3.5-9B Demoted Heads — MI300X 262K Context

Model Description

This is a hybrid attention variant of Qwen3.5-9B, trained with the RTPTurbo demoted-heads pipeline. The model uses a mixed linear attention (GDN) + full attention architecture to enable efficient long-context inference while maintaining retrieval quality.

Architecture

• 32 layers: 24 GDN (linear) + 8 full attention
• GQA: 16:4 (128 Q-heads, 32 KV-groups)
• Head dim: 256
• Max position: 262,144 tokens
• Demoted config: 1 mixed layer (L3: 3 local + 1 retrieval KV-group), 7 all-retrieval layers
• Local Q-heads: 12/128
• Window size: 8192
• Sinks: 4
• Indexers: 96

Training

• Stage 1: Train low-dimensional projection weights ($W_h^Q, W_h^K$) for retrieval heads only. Backbone frozen. Objective: KL divergence between full and projected attention distributions. AdamW (lr=1e-3, betas=(0.9, 0.999), wd=0.01, fused=True). ~840K trainable params.
• Stage 2: End-to-end self-distillation. Projections frozen, full backbone trained to match dense teacher's top-10 logits. $\mathcal{L}{distill} = \text{KL}(\text{softmax}(z{dense}^{(10)}) | \text{softmax}(z_{sparse}^{(10)}))$. AdamW (lr=3e-6, betas=(0.9, 0.95), wd=0.01, fused=True). ~600 steps at 48K avg seq len.

Inference (vLLM)

Served on AMD MI300X (gfx942, ROCm 6.16.13) with vLLM v0.23.1rc1.dev88.

Serve Command

python -m vllm.entrypoints.openai.api_server \
  --model /path/to/checkpoint \
  --max-num-seqs 100 \
  --max-model-len 262144 \
  --block-size 512 \
  --attention-backend ROCM_AITER_UNIFIED_ATTN \
  --compilation-config '{"cudagraph_mode":"FULL","splitting_ops":[]}' \
  --served-model-name qwen3.5-9b-demoted

Environment Variables

VLLM_ROCM_USE_AITER=1
VLLM_ROCM_USE_AITER_UNIFIED_ATTENTION=1
RTPURBO_LOCAL_ROCM_AITER=1

Performance

• Decode throughput: 137–157 tok/s at 8K context, 127 tok/s at 64K context
• KV cache: 155.98 GiB, 5.5M tokens capacity, 21x max concurrency at 262K
• FULL cudagraph mode enabled (local heads routed to RocmAiterUnifiedAttentionBackend)

Evaluation: AA-LCR (Long Context Retrieval)

Metric	Score
Overall	66/100 (66%)
Company Documents	71.4%
Academia	80.0%
Legal	66.7%
Government	54.5%
Industry	50.0%
Marketing	50.0%

• Judge: Qwen3.6-27B (served on same MI300X)
• Generation: 100 concurrent requests, temperature=1.0, top-p=0.95, top-k=20, presence-penalty=1.5, max-new=8192
• Baseline comparison: matches 67/100 non-demoted baseline — no degradation from demotion

Files

• model.safetensors-0000{1,2,3,4}-of-00004.safetensors — model weights (~19 GB total)
• model.safetensors.index.json — weight index
• model-rtpurbo-indexers-of-1.safetensors — indexer parameters
• config.json — model configuration
• rtpurbo_probe.json, rtpurbo_probe_demoted.json, rtpurbo_probe_allretr_backup.json — probe configs
• tokenizer.json, tokenizer_config.json, vocab.json, merges.txt — tokenizer
• preprocessor_config.json
• aa_lcr_demoted_9b_mi300x.jsonl — AA-LCR raw generations
• aa_lcr_demoted_9b_mi300x_judged.jsonl — AA-LCR graded results