whn09 f07d0a7c6f Neuron: compiled log-softmax + greedy argmax in the sampler; eager op census (#27)
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>
2026-10-23 09:42:08 +08:00
..
__init__.py initial commit 2026-10-20 12:09:42 +08:00
base.py Neuron: full DeepSeek-V4.1-Flash load + decode region graphs at EP=64 (issues #21, #25) 2026-10-23 01:39:22 +08:00
cache_engine.py Neuron: separate state regions, chunk-wide indexer context (issue #26) 2026-10-23 03:46:56 +08:00
cuda_graph_runner.py initial commit 2026-10-20 12:09:42 +08:00
decode_perf_cache.py initial commit 2026-10-20 12:09:42 +08:00
dflash_cuda_graph_runner.py initial commit 2026-10-20 12:09:42 +08:00
model_runner.py Neuron: compiled log-softmax + greedy argmax in the sampler; eager op census (#27) 2026-10-23 09:42:08 +08:00
neuron_graph_runner.py Neuron: compiled log-softmax + greedy argmax in the sampler; eager op census (#27) 2026-10-23 09:42:08 +08:00
runtime_config.py initial commit 2026-10-20 12:09:42 +08:00
worker.py initial commit 2026-10-20 12:09:42 +08:00
zos_status.py initial commit 2026-10-20 12:09:42 +08:00