Pull requests / #197
#197 sampler: top_k by lexicographic threshold, split over the whole GPU
closed · @gputier · 0 comentarios · En GitHub
AMD / HIPNVIDIA / CUDAModels & quants
Descripción
First of the three packages from #149, rebased on 0.1.26 as asked. The grouped Q2_0 kernel and the IQ grid decode follow once #154 is through. The sampled path (penalties, top_k, top_p, min_p, temperature, draw) kept one row on one SM. `sampler_kernel` excluded the logits already kept by sweeping the kept list for every logit of every round, which is O(k^2 x n_vocab) shared-memory compares per row, and it ran the double-precision tail on all 1,024 threads. What changes: - Split top_k, the new default. Stage 1 cuts each row into 4,096-logit blocks (61 per row at 248,320 logits); each warp keeps the top_k of its 1,024 penalised logits in registers and the block merges its four warp lists. Stage 2, one warp per row, merges the block lists and runs the tail on one warp. The block lists live in a per-stream scratch buffer grown on demand. - `sampler_one_block_kernel`: one block per row, each round takes the logits strictly after the previous pick in the selection order (value descending, id ascending). It is the fallback when the split cannot run (stream under graph capture, vocabulary above 262,144, more than 64 rows, no scratch). `STRATA_SAMPLER_ONE_BLOCK=1` forces it. - `sampler_kernel` is unchanged and stays as the reference; `STRATA_OLD_SAMPLER=1` selects it. - `hip_compat/cuda_runtime.h` maps `cudaStreamIsCapturing`, `cudaStreamCaptureStatus` and `cudaStreamCaptureStatusNone`, which the new capture check needs, since `sampler.cu` is in the shared CUDA/HIP `strata_kernels` list. Default behaviour: the split path is new code, but the list, its order and the pick are the same by construction (the first k of a union lie within the first k of each part, and a merge of ordered lists is ordered). The tail keeps the same arithmetic, with exps and quotients shared across a warp and lane 0 running the ordered sums and scans. Parity test: `sampler_parity` (fixtures 16 to 18) pins the whole top_k list position by position under ties, with rows holding +inf, -inf, NaN, +0/-0, the id-0 sentinel and penalised tokens landing on other tokens' values, over vocabularies of 248,320, 100,003, 262,144, 262,145 and 1,000. It also checks tied draws through the whole chain and the automatic fallbacks (graph capture, row cap). Picks are token ids compared for equality with a host reference. ctest runs it once per path, `sampler_parity`, `sampler_parity_one_block` and `sampler_parity_old`, so all three are held to the same expected ids. What I could and could not check: the commit applies on 0.1.26 without conflicts. I could not compile CUDA or HIP, or run `sampler_parity` on a GPU, from where I am now; your gate is the check there, and I will add results from my RTX 5090 when I have them. On HIP, `__syncwarp` is used unmapped, as in `native_flash_attn.cu`; I have not built it with ROCm.
En el sitio
Enlaces a install, modelos, releases.