[Feature] Support Step 3.5 Flash (StepFun) Model #1367

Checklist

• 1. If the issue you raised is not a feature but a question, please raise a discussion at https://github.com/sgl-project/sgl-jax/discussions/new/choose Otherwise, it will be closed.
• 2. Please use English, otherwise it will be closed.

Motivation

Add support for Step 3.5 Flash, StepFun's open-source sparse-MoE LLM (released 2026-01, Apache 2.0). It is a strong "intelligence-density" agentic model (~196B total / ~11B active) and is architecturally close to MiniMax-M2 (already supported in sglang-jax), so most existing MoE building blocks can be reused with no new kernel and no new state structure (it is plain GQA softmax attention, not MLA / not linear attention).

Model configuration (verified against config.json + modeling_step3p5.py):

• Architecture: Step3p5ForCausalLM (model_type: step3p5), decoder-only sparse MoE
• ~196B total / ~11B active per token (incl. 0.81B MTP head, per technical report)
• 45 layers, hidden_size 4096, vocab 128896, max_position 262144 (256K), BF16
• Attention: GQA (att_impl_type: GQA), head_dim 128, per-head QK-norm (zero-centered / Gemma-style RMSNorm), head-wise attention gate
  • Heterogeneous layers (3:1 pattern, full at layers 0,4,8,…,44 → 12 full / 33 sliding):
    • full-attention layers: 64 query / 8 KV heads, partial_rotary_factor 0.5, rope_theta 5e6 + llama3 rope scaling (factor 2.0)
    • sliding-attention layers: 96 query / 8 KV heads, partial_rotary_factor 1.0, rope_theta 1e4, window 512, no rope scaling (yarn_only_types: ["full_attention"])
• MLP: first 3 layers dense (intermediate 11264); layers 3–44 are MoE
• MoE: 288 routed experts + 1 shared expert (intermediate 1280 each), top-8, sigmoid routing computed in fp32 + router bias + normalized top-k weights + routed-scaling 3.0
• Clamped SwiGLU on the last layers only (swiglu_limits, value 0.0 means no clamp); MTP head (3 nextn-predict layers)
• Official weights: BF16 (also FP8 / Int8 / Int4 / GGUF)

Reuse map to existing sglang-jax components (no new kernel):

• GQA softmax attention → RadixAttention; sliding window via sliding_window_size (cf. gemma2)
• per-head QK-norm + zero-centered RMSNorm → GemmaRMSNorm + per-head pattern (cf. qwen3, gemma2)
• per-layer RoPE (rope_theta / partial_rotary_factor / llama3 scaling differ per layer) → existing RotaryEmbedding / Llama3RotaryEmbedding (cf. embeddings.get_rope)
• head-wise attention gate → one Linear(hidden → num_heads) + sigmoid, broadcast-multiplied over heads before o_proj (a few lines)
• first-k dense + MoE layers → first_k_dense_replace pattern (cf. deepseek_v3)
• sigmoid routing + router bias + renorm + routed-scaling + shared expert → GateLogit / TopK / EPMoE / FusedEPMoE (cf. minimax_m2, qwen2_moe)
• clamped SwiGLU → SiLU-and-mul with an elementwise clamp branch (a few lines)

Key difficulties / things to get right (correctness-sensitive, not blockers):

1. MoE routing order must match exactly: scores are sigmoid in fp32; the router bias is only used to select top-k, while the gating weights are gathered from the bias-free sigmoid scores; then normalize top-k weights and multiply by routed_scaling_factor = 3.0. Need to confirm TopK(correction_bias=...) in sglang-jax reproduces this exact semantics (bias affects selection only, not weighting).
2. Per-layer RoPE wiring: full vs sliding layers differ in rope_theta, partial_rotary_factor, and whether llama3 scaling applies — each layer must instantiate RoPE with the correct parameters.
3. Head-wise gate layout: gate comes from the attention input hidden states (not the attention output); after sigmoid it multiplies the [tokens, num_heads, head_dim] output per head, before o_proj. Broadcast axis must line up with the sglang-jax attention output layout.
4. Heterogeneous head counts under TP: full=64 / sliding=96 query heads must be divisible by TP size; KV=8 needs replication when TP>8 (the existing kv_head_padding / KV-head replica logic in minimax_m2 can be reused).
5. Per-head QK-norm: norm is applied per head (over head_dim) after reshaping, using the zero-centered (weight+1) variant.

MTP (nextn-predict) layers can be skipped for the initial base-inference support.

Related resources

• Hugging Face (BF16): https://huggingface.co/stepfun-ai/Step-3.5-Flash
• Hugging Face (FP8): https://huggingface.co/stepfun-ai/Step-3.5-Flash-FP8
• Hugging Face (Base): https://huggingface.co/stepfun-ai/Step-3.5-Flash-Base
• ModelScope mirror: https://modelscope.cn/models/stepfun-ai/Step-3.5-Flash
• GitHub: https://github.com/stepfun-ai/Step-3.5-Flash
• Technical report (arXiv): https://arxiv.org/abs/2602.10604
• Upstream sglang reference implementation: python/sglang/srt/models/step3p5.py
• vLLM deployment recipe: https://docs.vllm.ai/projects/recipes/en/latest/StepFun/Step-3.5-Flash.html

Implementation strategy

Phased verification — model code is written once (full parameterized file); phases control what gets verified, not what gets built.

• EPMoE first: enables GMM == per-expert-loop invariant check; FusedEPMoE switch deferred to Step 3.5 Flash: FusedEPMoE performance backend switch #1391
• MHA pool first: simplest KV layout for alignment; SWA switched on after model correctness is confirmed (guard: output(full) == output(SWA))
• MTP skipped: layers 45/46/47 ignored at load time; speculative decoding deferred to Step 3.5 Flash: MTP (Multi-Token Prediction) speculative decoding #1392

Sub-issues

• Step 3.5 Flash: FusedEPMoE performance backend switch #1391 — FusedEPMoE performance backend switch
• Step 3.5 Flash: MTP (Multi-Token Prediction) speculative decoding #1392 — MTP (Multi-Token Prediction) speculative decoding