feat: add qwen3 tts topk gumbel sampler

[l-wave]

Summary

This PR adds a Qwen3-TTS oriented short-vocab fused sampler operator, qwen3_tts_topk_gumbel_sample, for TTS/code-token sampling workloads.

The new op handles standalone logits and pre-generated uniform noise tensors:

• logits: [batch_size, vocab_size]
• noise: [batch_size, vocab_size]
• dtype: float32
• layout: contiguous
• supported vocab size: vocab_size <= 2048
• supported top-k: 1 <= topk <= 64
• production Qwen3-TTS shape: vocab_size=2048, topk=50

It returns sampled token ids:

• output: [batch_size, 1]
• dtype: int64

It exposes:

• Python API: hpc.qwen3_tts_topk_gumbel_sample
• Torch op: torch.ops.hpc.qwen3_tts_topk_gumbel_sample
• CUDA implementation: src/sampler/qwen3_tts_topk_gumbel_sample.cu
• tests/benchmark: tests/test_sampler.py

Motivation

Existing hpc-ops sampler kernels are designed for more general large-vocab sampling flows. The generic fused_sampler supports features such as repetition penalty, temperature, softmax policy, top-k, top-p, Gumbel-max sampling, and penalty-mask writeback. It is a feature-rich large-vocab path, but it is not the most direct fit for Qwen3-TTS code-token sampling.

The Qwen3-TTS code predictor path uses a much narrower sampling pattern:

scaled = logits * inv_temperature
topk_vals, topk_idx = torch.topk(scaled, topk)
scores = topk_vals - torch.log(-torch.log(noise[topk_idx]))
token = topk_idx[argmax(scores)]

For short TTS vocabularies such as V <= 2048, this PyTorch expression introduces multiple CUDA launches and intermediate tensors. Reusing the generic sampler also brings machinery that this path does not need, such as large-vocab multi-block top-k, top-p, softmax policy, repetition penalty, and penalty-mask writeback.

This PR adds a dedicated short-vocab fused sampler for this TTS path.

What changed

This PR adds a fused CUDA kernel for short-vocab TTS top-k Gumbel sampling:

• fuses top-k candidate selection, Gumbel transform, and argmax into one CUDA launch
• consumes caller-provided uniform random noise, matching the production substitution boundary
• supports vocab sizes up to 2048
• supports topk up to 64
• includes a production-specialized topk=50 path
• uses one CUDA block per batch row
• materializes only the compact top-k candidate set
• avoids producing a full masked logits tensor
• keeps the implementation separate from the existing generic fused_sampler

The PR also adds minimal-build support used by focused operator validation:

• HPC_OPS_IMPORT_MODULES=sampler imports only sampler Python APIs
• missing version / built_json ops are tolerated only in explicit minimal-import mode
• test import path now supports both build/lib.* and in-place _C.abi3.so

Difference from existing sampler

The existing fused_sampler is a general sampler for large-vocab decoding. It supports a broad feature set:

• float32 / bfloat16 logits
• repetition penalty
• temperature
• softmax before or after top-k
• top-k / top-p
• Gumbel-max
• penalty-mask writeback
• optional internal RNG
• large vocabulary such as 120832

The new qwen3_tts_topk_gumbel_sample is intentionally narrower:

• only float32 logits
• only external uniform noise
• no top-p
• no softmax policy
• no repetition penalty
• no penalty-mask writeback
• no internal RNG
• short vocab only, V <= 2048
• topk <= 64

This specialization lets the kernel avoid generic large-vocab sampler overhead and focus on the Qwen3-TTS code-token hot path.

Test coverage

Correctness is covered with a general TTS short-vocab matrix:

• vocab sizes: 64, 128, 256, 512, 1024, 1536, 2048
• top-k values: 1, 4, 16, 32, 50, 64
• batch sizes: 1, 4, 8, 16, 32, 64
• multiple inv_temperature values
• total golden correctness cases: 42

It also covers error and boundary cases:

• invalid topk
• vocab_size > 2048
• unsupported logits dtype
• unsupported noise dtype
• noise shape mismatch
• non-contiguous logits
• non-contiguous noise

Correctness command:

HPC_OPS_IMPORT_MODULES=sampler pytest tests/test_sampler.py -q -k "qwen3_tts_topk_gumbel_sample and not benchmark"

Current result:

52 passed, 65 deselected

Benchmark

Environment:

• GPU arch: SM90 / H20
• dtype: FP32 logits and FP32 uniform noise
• benchmark boundary: post-noise sampler chain only
• benchmark excludes random noise generation from both HPC and PyTorch paths
• benchmark command:

HPC_OPS_IMPORT_MODULES=sampler \
HPC_OPS_RUN_BENCHMARKS=1 \
HPC_OPS_BENCH_WARMUP=200 \
HPC_OPS_BENCH_ITERS=2000 \
pytest tests/test_sampler.py -q -s -k qwen3_tts_topk_gumbel_sample_benchmark

Benchmark result summary, excluding the removed topk=1 micro-case:

shape	batch	vocab	topk	torch us	hpc us	speedup
vocab64_top64	8	64	64	59.796	30.936	1.93x
vocab256_top16	8	256	16	63.871	10.229	6.24x
vocab512_top16	1	512	16	63.259	10.679	5.92x
vocab512_top50	16	512	50	60.008	16.311	3.68x
vocab1024_top32	8	1024	32	64.171	18.845	3.41x
vocab1024_top64	32	1024	64	60.100	32.887	1.83x
vocab1536_top50	16	1536	50	60.048	18.729	3.21x
vocab2048_top16	8	2048	16	63.633	14.461	4.40x
vocab2048_top50	1	2048	50	59.300	19.571	3.03x
vocab2048_top50	8	2048	50	60.332	19.879	3.03x
vocab2048_top50	32	2048	50	60.148	20.078	3.00x
vocab2048_top64	64	2048	64	59.966	35.720	1.68x

Across the representative short-vocab TTS matrix, the fused op is faster than the PyTorch reference, with around 1.68x–6.24x speedup. For the production Qwen3-TTS setting V=2048, topk=50, the fused op is stable around ~20us and about 3.0x faster than the PyTorch reference across B=1/8/32.

Notes

This operator is intentionally narrower than a fully generic sampler. It currently targets:

• short-vocab TTS/code-token sampling
• float32 logits
• external uniform noise
• vocab_size <= 2048
• topk <= 64
• Qwen3-TTS production-style topk=50

It does not replace the existing generic fused_sampler; it complements it with a specialized fast path for Qwen3-TTS short-vocab sampling.