002f9879b2
Replaces cherry-picked upstream_ref with complete source trees. xllm/ — Iluvatar official C++ inference engine (15MB, 1470 files) Complete: kernels → layers → models → runtime → scheduler → api Excluded: .git, binary images, third_party submodule checkouts ds_vllm/ — Iluvatar official vllm fork (8MB, 703 files) Included: csrc/ (ALL CUDA kernels), fused_moe/, qwen3_5 model, _custom_ops Excluded: tests, benchmarks, docs, examples (not needed for reference) Critical call chains now fully traceable: MoE: moe_topk_softmax_kernels.cuh → ixformer.h → fused_moe.cpp → layer GDN: qwen3_gated_delta_net_base.cpp → qwen3_5_gated_delta_net.cpp Attention: ixformer.h → xllm_paged_attention → attention.cpp
1.4 KiB
Topk & Topp Operator Optimization
Background
In natural language generation tasks, the topK and topP sampling strategies are widely used to control the diversity and quality of generated text. However, in small models, the computation time for these two strategies is relatively long. This is mainly due to the fewer parameters in small models, which leads to reduced efficiency in sorting and filtering when processing probability distributions, thereby affecting generation speed. Therefore, optimizing the implementation of topK and topP in small models can enhance their sampling efficiency.
Feature Introduction
The implementation of the topKtopP operator merges multiple small operators, such as sorting, topK, softmax, and topP, into a single large operator, thereby improving computational efficiency and performance.
User Interface
Operator Call API
void top_k_top_p(torch::Tensor& logits,
const torch::Tensor& topK,
const torch::Tensor& topP);
logits: The input logits tensor containing the model's output scores.topK: The threshold tensor for selecting the top K probabilities.topP: The threshold tensor for selecting the cumulative probabilities.
Performance Effect
- After using the topKtopP fused operator, in the qwen2-0.5B model, TTOT decreased by 37%, and TTFT increased by 10%.