feat: support ling-2.6-flash

[cjx0709]

Summary

This PR adds Data Parallelism (DP) support to LightningAttnBackend for GLA (Gated Linear Attention) and includes a comprehensive test suite, plus end-to-end validation on Ling-2.6-Flash (BailingMoeHybrid) over TPU v6e-16.

Key Changes

• LightningAttnBackend: DP-aware state management with shard_map

  • Support DP>1 with P("data", "tensor", None, None) sharding for recurrent state
  • DP-local indexing for state gather/scatter operations
  • Varlen format with DP-sharded cu_seqlens in extend mode
  • Per-layer slope storage in tp_slope dict

• LinearRecurrentAttnBackend: DP metadata computation

  • cu_q_lens: 1D array [dp_size * (per_dp_bs_size+1)] for DP>1
  • recurrent_indices: Sharded along P("data") axis

• Test Coverage: 32 tests total across 4 files

  • 7 DP correctness tests (new file: test_lightning_backend_dp.py)
  • 18 end-to-end integration tests
  • 5 backend unit tests
  • 2 model-level tests

• Cleanup: Removed torch dependency tests (HF parity tests)

  • JAX reference implementation provides sufficient validation

• Eval coverage: added test/srt/eval/simple_eval_{aime26,csimpleqa}.py (driven by the existing test/srt/run_eval.py runner) for the validation below.

Testing

All 32 tests passing on TPU v6e-4.

Documentation

See updated RFC in docs/design/bailing_moe_linear_attention.md for detailed DP implementation and test coverage.

Validation on Ling-2.6-Flash (BailingMoeHybrid, TPU v6e-16)

End-to-end accuracy on the inclusionAI/Ling-2.6-flash model, served by this branch. Numbers compared to the official table on https://www.modelscope.cn/models/inclusionAI/Ling-2.6-flash.

Server environment

• Hardware: TPU v6e-16, 4 nodes × 4 chips (nnodes=4, one rank per node).
• Model: /models/Ling-2.6-flash (BailingMoeV2.5 hybrid attention).
• Repo path on each node: /sglang-jax, branch cjx/gla-kernel.
• Python: /opt/venv/bin/python (sgl_jax installed).

Launch command

Run the following on each rank with NODE_RANK ∈ {0,1,2,3} and MASTER_ADDR set to rank-0's reachable address. All four ranks should be started concurrently (do not wait for any rank to become ready before starting the next).

PYTHONPATH=/sglang-jax/python \
/opt/venv/bin/python -u -m sgl_jax.launch_server \
  --model-path /models/Ling-2.6-flash \
  --trust-remote-code \
  --tp-size 16 --dp-size 4 --ep-size 16 \
  --moe-backend epmoe \
  --nnodes 4 --node-rank "${NODE_RANK}" \
  --dist-init-addr "${MASTER_ADDR}:10011" \
  --host 0.0.0.0 --port 30271 \
  --page-size 128 \
  --context-length 65536 \
  --chunked-prefill-size 2048 \
  --dtype bfloat16 \
  --mem-fraction-static 0.90 \
  --max-running-requests 256 \
  --attention-backend fa \
  --disable-radix-cache \
  --skip-server-warmup \
  --log-level info

Wait for 'The server is fired up and ready to roll!' in the log on rank 0 (typical: ~5–6 min).

Eval results (vs official)

Category	Benchmark	Ling-2.6-flash (official)	This branch	Δ
Knowledge	C-SimpleQA (judge_accuracy, same-endpoint judge)	60.23	60.87	+0.64
Math	AIME26 (30 problems)	73.85	76.67	+2.82

Notes:

• C-SimpleQA uses the official LivingFutureLab/ChineseSimpleQA GRADER_TEMPLATE; judge runs against the same served endpoint (biased vs GPT-4o, but lands within 1pp here).
• AIME26 with 30 problems has high per-question variance; +2.82pp ≈ < 1 problem of difference and is within sampling noise.
• Both runs: 0 request errors / 0 truncations.

Eval commands

Run from any node that can reach the served port (here we hit the local rank-0 server at 127.0.0.1:30271). Both evals download their dataset from HuggingFace on first run.

AIME26 (full 30 problems, parallel=30, max_tokens=32768 — requires --context-length 65536)

cd /sglang-jax && \
/opt/venv/bin/python test/srt/run_eval.py \
  --base-url http://127.0.0.1:30271 \
  --eval-name aime26 \
  --num-threads 30 \
  --temperature 0.6 \
  --max-tokens 32768

Chinese-SimpleQA (full 3000 questions, parallel=256, with same-endpoint judge)

cd /sglang-jax && \
/opt/venv/bin/python test/srt/run_eval.py \
  --base-url http://127.0.0.1:30271 \
  --eval-name csimpleqa \
  --num-threads 256 \
  --temperature 0.0 \
  --max-tokens 1024

Closes #1032

[gemini-code-assist]

Warning

You have reached your daily quota limit. Please wait up to 24 hours and I will start processing your requests again!

[JamesBrianD]

Apologies—let's wait until Kimi linear has been merged in first before merging this.

Apologies—let's wait until Kimi linear has been merged in first before merging this.