feat: add qwen3 tts rope operator

[l-wave]

Summary

This PR adds a Qwen3-TTS oriented NeoX RoPE operator, qwen3_tts_rope, for short-sequence TTS/code-group workloads.

The new op handles standalone Q/K tensors:

• q: [B, num_q_heads, seq_len, head_dim]
• k: [B, num_kv_heads, seq_len, head_dim]
• cos/sin: [B, seq_len, head_dim]
• dtype: bfloat16
• currently specialized for head_dim=128
• layout: NeoX-style rotate_half
• supported layout: leading dimensions may be strided, but the last dimension must be contiguous

It returns contiguous rotated Q/K tensors and exposes:

• Python API: hpc.qwen3_tts_rope
• Torch op: torch.ops.hpc.qwen3_tts_rope
• CUDA implementation: src/rope/qwen3_tts_rope.cu
• tests/benchmark: tests/test_tts_rope.py

Motivation

Existing hpc-ops RoPE kernels are optimized for attention/KV-cache paths, especially packed QKV + KV-cache store flows such as rope_norm_store_kv and rope_norm_store_kv_fp8.

Those kernels are not a direct fit for TTS short-sequence code-predictor style workloads where Q and K are already materialized as separate tensors and the caller only needs NeoX RoPE applied to standalone Q/K tensors.

For this TTS path, the PyTorch expression:

q_out = (q * cos) + (rotate_half(q) * sin)
k_out = (k * cos) + (rotate_half(k) * sin)

introduces multiple elementwise ops, intermediate tensors, and launch/dispatch overhead. This overhead is especially visible for short sequence lengths such as S=8/16/17/32.

What changed

This PR adds a dedicated fused CUDA kernel for TTS NeoX RoPE:

• fuses Q and K RoPE into a single CUDA launch
• supports GQA/MQA/no-GQA style head layouts via separate num_q_heads and num_kv_heads
• supports strided leading dimensions for Q/K/cos/sin tensors
• requires the last dimension of Q/K/cos/sin to be contiguous
• writes contiguous Q/K outputs
• uses row-packed CTA scheduling:
  • bs <= 1: 64 threads/row x 4 rows/CTA
  • bs > 1: 32 threads/row x 8 rows/CTA
• uses a BF16x2 fast path for aligned bs > 1 inputs to reduce global load/store and conversion overhead
• broadcasts per-row metadata inside the BF16x2 path to avoid repeating row index decomposition across all lanes
• keeps a scalar fallback path for unsupported alignment/stride cases
• keeps the implementation separate from existing KV-cache RoPE kernels

The PR also adds explicit input validation for the supported layout:

q.stride(3) == 1
k.stride(3) == 1
cos.stride(2) == 1
sin.stride(2) == 1

The test/benchmark matrix covers representative TTS-like shapes:

• qwen3_tts: S=16, Hq=16, Hkv=8
• qwen3_tts_full_groups: S=17, Hq=16, Hkv=8
• small_tts: S=8, Hq=8, Hkv=4
• gqa1: S=17, Hq=16, Hkv=1
• no_gqa: S=16, Hq=8, Hkv=8
• wide_heads: S=32, Hq=32, Hkv=8

Benchmark

Environment:

• GPU arch: SM90 / H20
• dtype: BF16
• op shape: standalone Q/K NeoX RoPE
• benchmark command:

HPC_OPS_ROPE_ONLY_LOAD=1 python3 tests/test_tts_rope.py --bench --iters 500 --warmup 50

Correctness:

HPC_OPS_ROPE_ONLY_LOAD=1 pytest tests/test_tts_rope.py -q
# 20 passed

Benchmark result summary:

shape	batch	rows	eager us	torch.compile us	triton us	hpc us	compile / hpc	triton / hpc
qwen3_tts	1	384	66.45	32.66	23.05	6.83	4.78x	3.38x
qwen3_tts	8	3072	66.75	32.40	23.16	6.98	4.64x	3.32x
qwen3_tts	32	12288	66.44	32.04	23.37	6.86	4.67x	3.40x
qwen3_tts_full_groups	1	408	66.60	33.18	23.08	6.77	4.90x	3.41x
qwen3_tts_full_groups	8	3264	66.51	32.25	23.37	6.97	4.63x	3.35x
qwen3_tts_full_groups	32	13056	65.93	32.19	23.42	6.86	4.69x	3.41x
small_tts	1	96	66.82	32.61	23.31	6.79	4.80x	3.43x
small_tts	8	768	66.55	32.28	23.25	6.79	4.75x	3.43x
small_tts	32	3072	66.64	32.38	23.18	6.95	4.66x	3.33x
gqa1	1	289	66.63	32.81	23.27	6.84	4.80x	3.40x
gqa1	8	2312	66.35	32.48	23.40	6.80	4.78x	3.44x
gqa1	32	9248	66.38	32.70	23.30	6.88	4.75x	3.38x
no_gqa	1	256	66.22	26.60	23.22	6.76	3.93x	3.43x
no_gqa	8	2048	66.16	26.38	23.21	6.78	3.89x	3.42x
no_gqa	32	8192	66.50	27.07	23.17	6.87	3.94x	3.37x
wide_heads	1	1280	66.46	32.36	23.46	6.84	4.73x	3.43x
wide_heads	8	10240	66.97	32.44	23.37	6.86	4.73x	3.41x
wide_heads	32	40960	114.26	32.47	28.53	16.46	1.97x	1.73x

Across this representative TTS short-sequence matrix, the fused CUDA op is consistently faster than both the compiled PyTorch reference and the Triton reference.

Compared with torch.compile, the speedup ranges from about 1.97x to 4.90x. Compared with the Triton reference, the speedup ranges from about 1.73x to 3.44x.

For the largest tested workload, wide_heads B=32, the BF16x2 path improves the standalone RoPE kernel from roughly 16.52us to 16.46us in this run while preserving the short-shape latency around the 6.8–7.0us launch-dominated region.

Notes

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

• NeoX rotate_half layout
• head_dim=128
• BF16 Q/K/cos/sin
• standalone Q/K TTS short-sequence workloads
• leading-dimension strided inputs with contiguous last dimension

It does not replace the existing packed-QKV + KV-cache RoPE kernels.