Custom topk-logprobs kernel + Remove a redundant Device to Host logits copy.

[pramodith]

What does the PR do?

The extraction of top-k vocab indices from the logit scores of the draft model currently happens on the cpu. This involves a costly transfer of the entire logit scores from the GPU to CPU i.e. Device to Host. In this PR we move the top-k computation to the GPU entirely.

This PR also adds a verify-argmax fastpath that skips the
logits transfer from Device to Host and re-uses the argmax indices stored as an attribute. It can be toggled via DFLASH_GPU_VERIFY_ARGMAX. It re-uses the same CPU fallback on any out-of-range index.

A comment stated that previous builds had issues with GPU based verify-argmax, this PR adds a check to detect failures and resort to using the prior CPU method for running the top-k computation in this scenario.

The GPU kernel and verify-argmax is enabled by default but can be disabled via the DFLASH_GPU_DRAFT_TOPK and DFLASH_GPU_VERIFY_ARGMAX flags.

Files edited

• Minor bug fixes to bench_llm.py to use the right tokenizer + correct the naming conventions of datasets being used.
• draft_topk_cuda.cu/h: Creates a new cuda kernel used by the draft model to identify the topk (where k is in the range 1-8) vocab indices from the logit scores.
• qwen35_dflash_target.cpp: Invokes the new .cu kernel and falls back to the cpu execution of topk incase there are errors.
• test_dflash.cpp: Invokes the .cu kernel if the DFLASH_GPU_DRAFT_TOPK flag is enabled. Uses the fast-path verification of draft tokens if DFLASH_GPU_VERIFY_ARGMAX=1, if it is set to 2 it identifies/reports any mismatches between cpu and gpu based verification.
• test_draft_topk_cuda.cpp: Correctness tests for the cuda kernel by comparing it against the previously used cpu function.
• CMakeLists.txt: Register the new cuda kernel.

Results

Results in an ~27% increase in tok/s. (105.63-82.80)/(82.8)
Numbers can be reproduced by running bench_llm.py after re-building with the changes.
The benchmark was run on an RTX 3090.

Results before

[bench] ==== HumanEval (n=10, n_gen=256) ====
  [01/10] n_tok=  92  AR= 33.96  DFlash=  98.69  AL= 6.74
  [02/10] n_tok= 146  AR= 33.93  DFlash=  76.72  AL= 5.33
  [03/10] n_tok= 142  AR= 33.55  DFlash=  99.51  AL= 7.11
  [04/10] n_tok= 128  AR= 34.39  DFlash=  75.50  AL= 5.57
  [05/10] n_tok= 180  AR= 33.09  DFlash=  87.92  AL= 6.74
  [06/10] n_tok= 126  AR= 33.31  DFlash=  71.80  AL= 5.45
  [07/10] n_tok=  59  AR= 33.48  DFlash=  71.72  AL= 5.45
  [08/10] n_tok= 149  AR= 33.33  DFlash=  75.89  AL= 5.69
  [09/10] n_tok= 133  AR= 33.19  DFlash=  80.56  AL= 6.10
  [10/10] n_tok= 103  AR= 33.10  DFlash=  89.65  AL= 6.56
  HumanEval mean: AR=33.53  DFlash=82.80  AL=6.07  2.47x

[bench] === SUMMARY ===
Task                AR    DFlash      AL   Speedup     Score
HumanEval        33.53     82.80    6.07     2.47x          

Results after

[bench] ==== HumanEval (n=10, n_gen=256) ====
  [01/10] n_tok=  92  AR= 33.66  DFlash= 118.42  AL= 6.74
  [02/10] n_tok= 146  AR= 34.09  DFlash=  93.11  AL= 5.33
  [03/10] n_tok= 142  AR= 34.33  DFlash= 123.95  AL= 7.11
  [04/10] n_tok= 128  AR= 33.73  DFlash=  97.60  AL= 5.57
  [05/10] n_tok= 180  AR= 34.15  DFlash= 117.04  AL= 6.74
  [06/10] n_tok= 126  AR= 34.01  DFlash=  94.94  AL= 5.45
  [07/10] n_tok=  59  AR= 34.60  DFlash=  94.64  AL= 5.45
  [08/10] n_tok= 149  AR= 33.02  DFlash=  97.45  AL= 5.69
  [09/10] n_tok= 133  AR= 33.90  DFlash= 104.93  AL= 6.10
  [10/10] n_tok= 103  AR= 34.06  DFlash= 114.25  AL= 6.56
  HumanEval mean: AR=33.95  DFlash=105.63  AL=6.07  3.11x

[bench] === SUMMARY ===
Task                AR    DFlash      AL   Speedup     Score
HumanEval        33.95    105.63    6.07     3.11x          

Reproducing Results

Baseline: DFLASH_GPU_VERIFY_ARGMAX=0 DFLASH_GPU_VERIFY_ARGMAX=0 python server/scripts/bench_llm.py --bench HumanEval

Latest Results: DFLASH_GPU_VERIFY_ARGMAX=1 DFLASH_GPU_VERIFY_ARGMAX=1 python server/scripts/bench_llm.py --bench HumanEval

[DeanoC]

Thanks for the extra context, LGTM then