coffee/DeepInfer
1
0
Fork 0
mirror of https://github.com/deepseek-ai/DeepInfer.git synced 2026-10-29 07:17:08 +08:00
Code Issues Packages Projects Releases Wiki Activity
DeepInfer/deepinfer/model_executor/layers/tilelang_ops/neuron
History
whn09 f4144fc3fc Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 
- vendor/deep_ep: EPBuffer dispatch/combine with torch.distributed uneven
  all_to_all (host-synced split sizes), EngramBuffer sharded lookup,
  BucketBuffer all_gather (CP-only reduce paths stubbed); verified against
  the DeepEP-Ascend references at world size 1/2/4.
- tilelang_ops/neuron sampling: random_sampling (asc two-level fixed-point
  sampler, Philox), fused_sample (logprobs / do_sample), sampler metadata,
  sum_smaller_probs, structured-output bitmask; DSpark/MTP/varlen kernels
  raise NotImplementedError. All 42 asc *_tl factories now have a Neuron
  counterpart.
- vendor/flash_attn: SDPA varlen stand-in for the ViT (mirrors npu_ops).
- transfer_to_neuron: keep `with torch.device(...)` working on patched
  factories; fix manual_seed recursion in host mode.
- engram: gather the packed ue8m0 sf as bytes on Neuron.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
2026-10-20 18:57:27 +08:00
..
__init__.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
_unsupported.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
candidate_blocks.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
compress.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
compute_seqlens_k_zos.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
dflash_copy.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
dflash.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
flash_mla_utils.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
fused_sample.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
fused_transform.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
gather_engram_tokens_zos.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
gather_input_ids_zos.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
gather.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
get_kv_index.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
image_embed_ops.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
kv_index_cat.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
kv_quant.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
load_paged_kv_cache.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
norm.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
prepare_cp_mirror.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
prepare_v4_decode_metadata.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
random_sampling.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
rstd_1ulp.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
sampler.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
schedule_mtp_varlen.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
state_cache.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
store_kv_cache.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00
store_mtp_info.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
structure_output.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
sum_smaller_probs.py Neuron basic mode — part 2: deep_ep, sampling kernels, fixes 2026-10-20 18:57:27 +08:00
swa_index.py Add AWS Trainium (Neuron) platform, basic mode — part 1 2026-10-20 18:34:51 +08:00