On-device greedy sampling (-Dllama.deviceSample): GPU argmax, keep logits device-side (~500x less D2H)

[mikepapadim]

On-device greedy sampling (-Dllama.deviceSample=true)

Appends a GPU argmax over the logits to the FP16 logits task graph, so on the greedy path only the sampled token id (1 int) crosses to the host each step — the full vocab logits row stays device-side (no D2H copy, no host CPU scan). This is the single-token analog of the batched-decode on-device sampling.

How

• TransformerComputeKernels.argmaxLogits — single-workgroup tree-reduce argmax over vocab, writes state.sampledToken[0].
• LogitsFP16Layer (gated on -Dllama.deviceSample): adds the argmax task, transfers only sampledToken, and skips the wrapLogits D2H.
• GPU decode loops (Llama/Mistral, Qwen3) read state.sampledToken instead of running the host sampler.
• Guarded in LlamaApp: enabled only for GPU + greedy (temperature 0) + FP16 + Llama/Mistral/Qwen3. Any other config (temperature > 0, non-FP16, other arch) clears the flag so the host still receives the full logits. Opt-in, default off — zero change to existing runs.

Verified — mechanism works (nsys, RTX 4090, Llama-1B, 60 tokens)

	Device→Host copy total	per token
host sampling (default)	13.34 MB (26 × 0.513 MB logits row)	513 KB
on-device sampling	0.001 MB (26 × 4 B token id)	4 B

~500× less D2H traffic, output bit-identical & coherent ("Paris, the capital of France…").

Honest throughput finding

At n=1 single-token decode this is throughput-neutral (Llama-1B 98.4→98.1 t/s; +graphs 110→110; Qwen3-1.7B 55.4→54.4): the eliminated copy is only one ~0.5 MB vocab row and overlaps the ~11 ms forward, so it isn't the bottleneck — the same n=1 launch/host-bound story as the hybrid-libs PR (#131). Where it pays off is batched decode (65 MB D2H → +30%, already shipped in #129) and serving/concurrency (removes per-step host CPU argmax, freeing the host thread). Shipping it here gives the single-request path the same device-resident-logits mechanism, correctly guarded.

@orionpapadakis — standalone feature off feat/mma_cuda; review welcome. Can extend to the Q8 / Granite logits layers if you want the mechanism everywhere.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.

[mikepapadim]

@orionpapadakis standalone on-device greedy sampling off feat/mma_cuda — GPU argmax, only the token id leaves the device (nsys: 13.3 MB → 1 KB D2H over 26 tokens). Opt-in -Dllama.deviceSample=true, guarded to GPU+greedy+FP16+Llama/Mistral/Qwen3. Throughput-neutral at n=1 (small logits row), the win is batched/serving — details in the PR body.