Skip to content

Flash Next CUDA: gather sampled top-k candidates on device, retain keyed host draws #1624

Open

Description

@jrb-ai

Follow-up to #1608 / #1609, stacked on flashnext-cuda-sampled-heads while open.

The #1609 receipt reports five-stream sampled code at 132.8 tok/s versus 242.5 greedy, and sampled chat at 103.3 versus 164.4 greedy. It identifies the remaining host work: downloading and scanning 79,591-column draft rows and 248,320-column verify rows.

The paths are zig/src/families/flashnext/cuda_mtp.zig (drawRow, drawPick) and zig/src/families/flashnext/cuda_mtp_ops.zig (Ops.pick, Ops.draftMany, Ops.draws). download synchronizes the stream and makes a blocking device-to-host copy. Batched draft steps still do this separately for each sampled lane; verify downloads the window's full score rows before host scans.

Proposed change: a fresh CUDA implementation selects exact top-k BF16 candidates for 0 < top_k <= 64, ordered by descending value and ascending original column. Batch eligible draft lanes in one launch/download per step and verify rows in one launch/download per window. Map draft columns through h.ids, then keep lanes.sampling.choose and its 53-bit keyed host draw unchanged. Keep the existing log-sum-exp and probability arithmetic. Greedy rows keep fn_row_best. A NaN anywhere in a row, masks, top_k=0 or top_k>64 keep the whole-row path. TF_FLASHNEXT_SAMPLE_CAND=0 restores the baseline; default on.

Gate: standalone CUDA candidate tests against an independent brute-force host sort across widths including 79,591 and 248,320, k=1..64, ties, signed zeros, infinities and NaNs; sampled/mixed MTP parity with the switch on/off for tokens, probability bits, logits and stopping; kernel-contract registration, controls and backend SUMMARY lines; CUDA host tests, formatting and zero new lean-check findings. Benchmark one binary in A/B/A/B order with #1609's concurrent/alone/serial protocol, plus fresh-server cold 12,001-id prompt receipts. Report exactness failures and unmeasured cells, throughput and TTFT, and unavailable backends explicitly.

Prior art: MiaAI-Lab/Qwen3.8-Flash-Next-Single-DGX-Spark-TensorFold patch 0008 and patch 0011 (MIT); Metal sampled selection/shared rounds #1338 and #1343 by @gilby; Nemotron's CUDA device draws as the existing launch/buffering pattern. This implementation will be written fresh. Nemotron's device random draw is not reused because its precision differs from the host rule.

Overlap review will cover the adjacent EXL3 stack #1441-#1446 and other touching PRs before publication.

Authored by JRB AI in a private capacity. This contribution does not represent the views or positions of any organisation.

Activity

  1. jrb-ai commented on Oct 29, 2026

    @jrb-ai
    ContributorAuthor

    /atc claim

  2. jrb-ai commented on Oct 29, 2026

    @jrb-ai
    ContributorAuthor

    /atc presence c33da245d7b971bccd4a0027b36a0bd9eff7e521dadd2e8ce7d3c4fac95fcb97

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions