stepfun-ai

/

Step-5-Preview

Like 137

Follow

StepFun 3.3k

We integrated Step 5 Preview into our own C++/CUDA inference engine (no PyTorch
at runtime) and converted the FP8 release in place to NVFP4: 606 GB of
safetensors → a 317 GB tree, with the 30 976 routed experts streamed from host
mmap through a device-slot LRU. The dense graph pins at 8.6 GiB; with a
~1 500-slot expert cache the process sits at roughly 23 GiB of VRAM on one
RTX 5090. Thanks for shipping the standalone inference/ implementation — the
packed-DSA reference made the hybrid schedule and the router arithmetic
unambiguous.

What worked well

• Greedy determinism is solid: the framed control prompt answers "Paris" in the
  reasoning block, and three runs produced byte-identical final logits
  (logits_hash = 63b0f8b2fcf56d8a over all 128 896 logits), including across a
  full model reload.
• The native tool wire is unambiguous: <tool_call> / <function=NAME> /
  <parameter=KEY> bodies parse cleanly. A real-weight round-trip generated the
  call after the </think> marker and parsed to {"city":"Paris"}.
• The router recipe is well specified: sigmoid scores, a selection-only bias,
  renormalised (score − bias), then × 3.0, with the gated shared expert
  added in FP32.
• The tokenizer bridge reproduces the checkpoint tokenizer exactly (7/7
  encode+decode cases, including the digit/CJK splits and the o200k
  punctuation-run pattern).

Three checkpoint facts that cost debugging time

1. The RMSNorm is zero-centred: x / rms(x) * (1 + w). The stored weights
   are centred on zero, so treating them as a conventional RMSNorm scale
   multiplies every activation by roughly zero. Shapes and dtypes do not object,
   and the first symptom is fluent-looking but empty output (in our case, a
   stream of newlines). A note in the model card would save any new integrator
   the afternoon it cost us.

2. moe_layer_list is 3..90, but layer 91 is a dense-MLP full-attention
   layer. Reading the config as "dense prefix then MoE to the end" makes the
   last layer look for 352 experts that do not exist. The mixed layout also
   means an expert store cannot assume a contiguous dense suffix; we bound the
   store at layer 91 and keep the flat index layer × experts + e.

3. Layers 88–90 ship BF16 experts. They are listed in
   quantization_config.modules_to_not_convert and have no
   weight_scale_inv. A converter that keys on "experts are FP8" silently stops
   at layer 88; it must accept both an FP8 block-scaled stack and a BF16 stack.

Things that were easy to get subtly wrong

• The expert projections are fused stacks [352, 1536, 4096] /
  [352, 4096, 1536], not per-expert tensors; the converter has to split them
  and re-quantise each expert.
• The RoPE span is per layer and the table is halved (pair (d, d + span/2),
  NEOX): 192 dims at theta 1e4 on the sliding layers, 64 dims at theta 1e7 on
  the full-attention layers. The rope_scaling llama3 block has factor = 1.0
  and is the identity.
• The per-head attention gate (self_attn.g_proj, one logit per head) is
  applied before o_proj.
• The final RMSNorm narrows the FP32 residual to BF16 first; that rounding
  boundary is part of the reference arithmetic.
• The router gate is BF16 on disk but must be evaluated in FP32.

What we did not run

The DSA indexer is stored in the converted tree but not executed. At or below
topk × region_block_size = 4096 tokens every region is selected, so the dense
causal path is the reference path; beyond that our engine refuses the request
rather than silently producing a different answer. Implementing the FP8 region
summaries, the radix top-k with its zero-group tie-break and the region gather
is the remaining work for the 1M-token claim.

One data point on speed

Our correctness-first path (token-by-token prefill/decode, a router readback per
MoE layer, no CUDA graphs, no prefetch stream) reaches 1.82 t/s at 2 048
generated tokens (TTFT ≈ 14.5 s) on the RTX 5090, with the 6.6 GB/token
expert stream as the bottleneck. That is a floor, not a verdict on the model's
throughput — it mostly measures our host-to-device staging.

Overall: a well-specified model to port, and the three checkpoint facts above
are the only things that cost real debugging time.