#38496 [Model Runner V2] Fuse probabilistic rejection sample kernels
原始 PR · 作者 TheEpicDolphin · 合并时间 2026-04-08 08:37
融合概率性拒绝采样内核,优化内存分配并消除softmax,提升推测解码性能。
建议核心工程师精读`probabilistic_rejection_sampler_utils.py`中的Triton内核实现,关注`_compute_block_max_and_sumexp`和`_probabilistic_rejection_kernel`的设计,以学习内核融合和数值稳定性优化技巧;同时,查看测试文件中的卡方检验方法,了解如何验证采样分布正确性。
参与讨论