perf(BACKEND-ROCM): widen sampling block to 1024 and split-phase random sample

[ghazni101]

Closes #3009.

What

Non-greedy decode on ROCm ran 22% slower than greedy because three per-row sampling kernels launched as <<<1, 256>>> — one block of 4 wavefronts on 1 CU out of 96, scanning the full ~152K vocab per token. This PR widens the per-row block to 1024 threads and adds a split-phase random sample kernel that spreads the Gumbel-score argmax across 128 blocks (all 96 CUs).

Why

rocprofv3 measured 3.34 ms/tok of sampling overhead at single-request decode:

Kernel	Before	After	Speedup
ApplyTopKTopPRowK	1270 µs	230 µs	5.5×
SoftmaxK	406 µs	106 µs	3.8×
RandomSampleK	1665 µs	~50 µs	33×
Total	3341 µs	~386 µs	8.6×

The RandomSampleK kernel is compute-bound on log(u) in double precision inside ExpNoise (the Gumbel-score noise source), so wider blocks alone barely helped (1665 → 1580 µs). The split-phase kernel is the critical fix: 128 blocks across all CUs, each thread handling 1–2 elements instead of 148.

How

Two changes in src/vt/rocm/rocm_sample.hip:

1. kVocabBlock = 1024 for per-row sampling kernels (top-k/top-p, softmax, random sample, min-p). 16 wavefronts vs 4 on the same CU. BlockRed* helpers use blockDim.x instead of the compile-time kBlock constant.

2. Split-phase random sample (VT_SAMPLE_SPLIT=1, default ON). Phase A (RandomSampleSplitAK) launches n × 128 blocks, each writing a per-block (score, index) partial to scratch. Phase B (RandomSampleSplitBK) reduces the 128 partials per row. Falls back to single-block when v < 4096 or n > 64. Scratch is a grow-only hipMallocAsync allocation (stream-ordered, same pattern as the greedy argmax split in rocm_dense_basic.hip).

ArgReduce is associative and order-independent (compares true global indices, not thread order), so the split-phase result is bit-identical to the single-block kernel for every input, ties included.

Measured result

Qwen3.5-4B Q4_K_M, fp8 KV cache, 256 tokens, 6 reps, warm median, gfx1100 (RX 7900 XTX, 96 CUs):

Path	Before	After	Change
Greedy	97.2 tok/s	97.8 tok/s	unchanged
Non-greedy	75.6 tok/s	96.5 tok/s	+27.6%
Gap	22%	3.4%

Correctness

• test_ops_sample: 29/29 pass (2 CUDA-only skipped), 278K assertions
• test_sampler: 21/21 pass
• A/B test: 20 seed/temperature combinations (seeds 1–12345, temps 0.3–1.5) produce bit-identical output with VT_SAMPLE_SPLIT=0 vs 1
• Determinism: 3 consecutive runs with same seed produce identical output
• Full test suite: 225/225 pass, 3 consecutive runs, zero failures
• Edge cases: 1-token generation, temp=2.0, top_k=1, temp=0.01 — all pass
• Distribution test (100K rows, V=4): passes — validates single-block fallback path

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:GLM-5-2 [TOOLS]

[ghazni101]

Rebased onto upstream/main (6f5e9dc). Clean rebase, no conflicts. Ping for re-review.

FOLLOWING_AGENTS_PROTOCOL

Following-Agents-Protocol: true
AI-Assisted: true
Assisted-by: AGENT:GLM-5-2 [OMP]