mirror of
https://github.com/deepseek-ai/DeepInfer.git
synced 2026-10-28 23:07:09 +08:00
The op census of a decode step (DEEPINFER_NEURON_REGION_TIMING) found 85 eager aten ops in the forward outside the region graphs but 1905 in sampling (the top-p/top-k bisection and the int64 philox draw), ~0.6 s of NEFF launches. - The log-softmax / top-p / top-k part runs as one compiled graph: 1.4 ms instead of 46.7 ms on one core for the [1, 4096] force_topk row, logprobs within 1e-9. - Greedy rows (temperature clamped to SAMPLING_EPS) have one-hot weights, so the two-level draw returns their argmax; an all-greedy batch takes the argmax from the graph and skips the draw. The philox draw stays eager for sampled rows: neuronx-cc graphs reject 64-bit constants outside the int32 range (NCC_ESFH001), an int32 philox is #3. Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com> |
||
|---|---|---|
| .. | ||
| tilelang_ops | ||
| __init__.py | ||
| attention_backend.py | ||
| layernorm.py | ||
| linear.py | ||
| mtp_info_pool.py | ||
| rotary_embedding.py | ||
| sampler_utils.py | ||
| sampler.py | ||
| vocab_parallel_embedding.py | ||