flashinfer.cake_sampling.choose_stage1

flashinfer.cake_sampling.choose_stage1(batch: int, vocab: int, sm_count: int | None = None, smem_limit: int | None = None, top_k_max: int | None = None, two_launch: bool | None = None) → tuple[int, int, bool]

(cluster, ept, stream) for batch rows of vocab entries (largest top-k top_k_max). The device-dependent defaults are resolved here and the ranking itself is memoised per resolved argument tuple (_choose_stage1_resolved): the served route asks twice per call and the round-7 fused-tail terms doubled the ranking’s cost (B200 host: 102 -> 162 us per call), which a host-bound two-launch chain pays inside its launch gap.

Register-resident candidates: fewest waves, then a register chunk of at least 16 entries, then the larger cluster. That resident choice is compared with every streaming variant through the fitted cost model of the device’s wave table (_STAGE1_COST_BY_SM_COUNT: resident waves * (base + per_ept * ept) against streaming waves * (base + chunk * chunks + cluster_cta * (cluster - 1)), both plus launch_cta * ctas / wave_ctas[1]); the resident variant wins ties and streaming ties prefer the smaller cluster. The wave table and its constants are selected by sm_count (the current device’s SM count when omitted; B200 148, H100 132 and Rubin R200 212 are measured, other counts use the nearest). top_k_max above _STAGE1_LARGE_K adds the table’s large-k streaming term (None ranks as a small top-k); two_launch says whether such a launch takes the stage-2/3 kernel separately (None: _launch_is_two_kernels on the current device, i.e. its whole-CTA tail policy) and selects the two-launch terms. Variants needing more dynamic shared memory than smem_limit (the current device’s opt-in limit when omitted) are not candidates.