{"paper":{"title":"FlashSampling: Fast and Memory-Efficient Exact Sampling","license":"http://arxiv.org/licenses/nonexclusive-distrib/1.0/","headline":"FlashSampling fuses exact categorical sampling into the LM-head matrix multiply so the full logits tensor is never written to HBM.","cross_cats":["cs.AI","cs.CL"],"primary_cat":"cs.LG","authors_text":"Mengdi Wang, Tomas Ruiz, Xuyang Shen, Yifan Zhang, Yiran Zhong, Zhen Qin","submitted_at":"2026-03-16T19:37:08Z","abstract_excerpt":"Sampling from a categorical distribution is mathematically simple, but in large-vocabulary decoding, it often triggers extra memory traffic and extra kernels after the LM head. We present FlashSampling, an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in HBM. The method is simple: compute logits tile-by-tile on chip, add Gumbel noise, keep only one maximizer per row and per vocabulary tile, and finish with a small reduction over tiles. In tensor-parallel decoding, FlashSampling replaces the all-gather of logits with streaming peer"},"claims":{"count":4,"items":[{"kind":"strongest_claim","text":"FlashSampling is an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in HBM; in tensor-parallel decoding it replaces all-gather with streaming peer-to-peer writes, achieving kernel-level speedups and up to 10% reduction in time per output token.","source":"verdict.strongest_claim","status":"machine_extracted","claim_id":"C1","attestation":"unclaimed"},{"kind":"weakest_assumption","text":"The claim that argmax decomposes exactly over vocabulary partitions (and that the hierarchical factorization for grouped variants preserves exact categorical sampling) holds without numerical error or edge-case failure when tiles are processed independently on-chip.","source":"verdict.weakest_assumption","status":"machine_extracted","claim_id":"C2","attestation":"unclaimed"},{"kind":"one_line_summary","text":"FlashSampling performs exact Gumbel-max sampling inside the LM-head matmul via on-chip tiling and hierarchical argmax reduction, delivering up to 10% faster token generation in vLLM on datacenter GPUs without approximation.","source":"verdict.one_line_summary","status":"machine_extracted","claim_id":"C3","attestation":"unclaimed"},{"kind":"headline","text":"FlashSampling fuses exact categorical sampling into the LM-head matrix multiply so the full logits tensor is never written to HBM.","source":"verdict.pith_extraction.headline","status":"machine_extracted","claim_id":"C4","attestation":"unclaimed"}],"snapshot_sha256":"ce3cba47136aac04afdbdcb5d0d61f83fc2176d9935ac901de94d3750e91ac7e"},"source":{"id":"2603.15854","kind":"arxiv","version":2},"verdict":{"id":"3f38b0a9-aca2-46ac-a984-5ab58db9a9e7","model_set":{"reader":"grok-4.3"},"created_at":"2026-05-15T09:57:45.197851Z","strongest_claim":"FlashSampling is an exact sampling primitive that fuses sampling into the LM-head matmul and never materializes the logits tensor in HBM; in tensor-parallel decoding it replaces all-gather with streaming peer-to-peer writes, achieving kernel-level speedups and up to 10% reduction in time per output token.","one_line_summary":"FlashSampling performs exact Gumbel-max sampling inside the LM-head matmul via on-chip tiling and hierarchical argmax reduction, delivering up to 10% faster token generation in vLLM on datacenter GPUs without approximation.","pipeline_version":"pith-pipeline@v0.9.0","weakest_assumption":"The claim that argmax decomposes exactly over vocabulary partitions (and that the hierarchical factorization for grouped variants preserves exact categorical sampling) holds without numerical error or edge-case failure when tiles are processed independently on-chip.","pith_extraction_headline":"FlashSampling fuses exact categorical sampling into the LM-head matrix multiply so the full logits tensor is never written to HBM."},"references":{"count":0,"sample":[],"resolved_work":0,"snapshot_sha256":"258153158e38e3291e3d48162225fcdb2d5a3ed65a07baac614ab91432fd4f57","internal_anchors":0},"formal_canon":{"evidence_count":0,"snapshot_sha256":"258153158e38e3291e3d48162225fcdb2d5a3ed65a07baac614ab91432fd4f57"},"author_claims":{"count":0,"strong_count":0,"snapshot_sha256":"258153158e38e3291e3d48162225fcdb2d5a3ed65a07baac614ab91432fd4f57"},"builder_version":"pith-number-builder-2026-05-17-v1"}