diff --git a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md index 98acba0ea8fe..8937a6477d00 100644 --- a/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md +++ b/docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md @@ -1,3 +1,8 @@ + + # Deployment Guide for Qwen3.8 MoE and Qwen3.5 MoE on TensorRT LLM - Blackwell Hardware ## Introduction @@ -146,6 +151,14 @@ The MTP draft length is configurable through `speculative_config.max_draft_len`; Each file assigns expert IDs to 544 global slots for every MoE layer. Repeated expert IDs create replicas of heavily routed experts. `layer_updates_per_iter: 0` keeps the assignment static. The no-MTP map covers the 92 target-model layers, while the MTP3 map also includes the MTP layer. These placements are specific to Qwen3.8-2.4T-A95B, EP32, the selected MTP mode, and the benchmark traffic distribution. Generate a new placement when any of those inputs changes. See the [EPLB example](https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/wide_ep/ep_load_balancer) for collecting routing statistics and generating a placement file. +### Experimental QKV Preprocessing Tiling + +`TRTLLM_ENABLE_QKV_ROPE_TILING=1` enables experimental eight-token tiling for fused Q/K normalization, rotary embedding, and gate extraction. It is disabled by default. Set the environment variable before starting `trtllm-serve` or importing TensorRT LLM in Python. + +Tiling applies to preprocessing batches with at least 256 tokens, at least 3,072 token/head pairs across Q/K/V (using local head counts after tensor parallelism), and `(head dimension, rotary dimension)` of `(128, 128)` or `(256, 64)`. Other workloads use the single-token path. Performance has been measured on B200 and Rubin; other GPU architectures can opt in, but their performance has not been established by those measurements. + +Tiling can change Q/K rounding. Changing batch size across a dispatch boundary can therefore change cached K values and generated tokens. Leave the flag unset when batch-invariant results are required. + ### Launch the TensorRT LLM Server Below is an example command to launch the TensorRT LLM server from within the container with the checkpoint that matches the configuration file you selected above. diff --git a/tensorrt_llm/_torch/attention/kernels/fused_qk_norm_rope_gate.py b/tensorrt_llm/_torch/attention/kernels/fused_qk_norm_rope_gate.py index 50d5e60d2c2d..847733bced10 100644 --- a/tensorrt_llm/_torch/attention/kernels/fused_qk_norm_rope_gate.py +++ b/tensorrt_llm/_torch/attention/kernels/fused_qk_norm_rope_gate.py @@ -23,12 +23,16 @@ LoRA, or attention-backend fallback behavior. """ +import os from typing import Optional, Tuple import torch import triton import triton.language as tl +_QKV_ROPE_TILING_ENABLED = os.environ.get("TRTLLM_ENABLE_QKV_ROPE_TILING", "0") == "1" +_MIN_TILED_HEAD_ROWS = 3072 # 384 eight-token CTAs. + @triton.jit def _fused_qkv_gemma_rmsnorm_rope_gate_kernel( @@ -58,108 +62,183 @@ def _fused_qkv_gemma_rmsnorm_rope_gate_kernel( use_mrope: tl.constexpr, mrope_section1: tl.constexpr, mrope_section2: tl.constexpr, + block_tokens: tl.constexpr, ): - token = tl.program_id(0).to(tl.int64) - head = tl.program_id(1) + """Normalize and rotate independent rows in a token/head tile. + + Normalized Q/K round through the output dtype before RoPE, as in the + single-token kernel. FP32 reduction order can differ across layouts. + """ q_size = num_q_heads * head_dim kv_size = num_kv_heads * head_dim output_dtype = tl.float16 if output_fp16 else tl.bfloat16 - + head = tl.program_id(1) + head_offsets = tl.arange(0, head_block) + rotary_offsets = tl.arange(0, rotary_block) + # Each program owns one Q, K, or V head across block_tokens tokens. + # Tiled tensors have shape [tokens, dimensions]; reductions stay per token. + if block_tokens == 1: + tokens = tl.program_id(0).to(tl.int64) + else: + tokens = (tl.program_id(0).to(tl.int64) * block_tokens + tl.arange(0, block_tokens))[ + :, None + ] + head_offsets = head_offsets[None, :] + rotary_offsets = rotary_offsets[None, :] + token_mask = tokens < num_tokens + tile_mask = token_mask & (head_offsets < head_dim) if head < num_q_heads + num_kv_heads: + # Input Q heads alternate with gates; output Q/K heads are packed. is_k = head >= num_q_heads - local_head = tl.where(is_k, head - num_q_heads, head) if is_k: - input_base = qkv_ptr + token * qkv_stride_token + 2 * q_size + local_head * head_dim - output_base = ( - qkv_out_ptr + token * qkv_out_stride_token + q_size + local_head * head_dim - ) + local_head = head - num_q_heads + input_offset = 2 * q_size + local_head * head_dim + output_offset = q_size + local_head * head_dim weight_ptr = k_weight_ptr else: - input_base = qkv_ptr + token * qkv_stride_token + local_head * 2 * head_dim - output_base = qkv_out_ptr + token * qkv_out_stride_token + local_head * head_dim + local_head = head + input_offset = local_head * 2 * head_dim + output_offset = local_head * head_dim weight_ptr = q_weight_ptr - - head_offsets = tl.arange(0, head_block) - head_mask = head_offsets < head_dim - x = tl.load(input_base + head_offsets, mask=head_mask, other=0.0).to(tl.float32) - weight = tl.load(weight_ptr + head_offsets, mask=head_mask, other=0.0).to(tl.float32) - inverse_rms = tl.rsqrt(tl.sum(x * x, axis=0) / head_dim + eps) - normalized = (x * inverse_rms * (weight + 1.0)).to(output_dtype).to(tl.float32) - - if has_pass_through: - pass_mask = head_mask & (head_offsets >= rotary_dim) - tl.store(output_base + head_offsets, normalized, mask=pass_mask) - - rotary_offsets = tl.arange(0, rotary_block) - rotary_mask = rotary_offsets < half_rotary - x_first = tl.load(input_base + rotary_offsets, mask=rotary_mask, other=0.0).to(tl.float32) - x_second = tl.load( - input_base + half_rotary + rotary_offsets, mask=rotary_mask, other=0.0 - ).to(tl.float32) - weight_first = tl.load(weight_ptr + rotary_offsets, mask=rotary_mask, other=0.0).to( + input_rows = qkv_ptr + tokens * qkv_stride_token + input_offset + output_rows = qkv_out_ptr + tokens * qkv_out_stride_token + output_offset + x = tl.load(input_rows + head_offsets, mask=tile_mask, other=0.0).to(tl.float32) + weight = tl.load(weight_ptr + head_offsets, mask=head_offsets < head_dim, other=0.0).to( tl.float32 ) - weight_second = tl.load( - weight_ptr + half_rotary + rotary_offsets, mask=rotary_mask, other=0.0 - ).to(tl.float32) - x_first = (x_first * inverse_rms * (weight_first + 1.0)).to(output_dtype).to(tl.float32) - x_second = (x_second * inverse_rms * (weight_second + 1.0)).to(output_dtype).to(tl.float32) + # Masked lanes contribute zero; RMS uses the true, unpadded head width. + inverse_rms = tl.rsqrt(tl.sum(x * x, axis=-1) / head_dim + eps) + if block_tokens != 1: + inverse_rms = inverse_rms[:, None] + # Gemma uses (1 + weight). Round to BF16/FP16 before FP32 RoPE math. + normalized = (x * inverse_rms * (weight + 1.0)).to(output_dtype).to(tl.float32) + if has_pass_through: + tl.store( + output_rows + head_offsets, + normalized, + mask=tile_mask & (head_offsets >= rotary_dim), + ) + + rotary_tile_mask = token_mask & (rotary_offsets < half_rotary) + if block_tokens == 1 or has_pass_through or (head_dim & (head_dim - 1)) != 0: + # Preserve scalar rounding and limit live registers for partial RoPE. + # Non-power-of-two heads also need loads that exclude padded lanes. + x_first = tl.load(input_rows + rotary_offsets, mask=rotary_tile_mask, other=0.0).to( + tl.float32 + ) + x_second = tl.load( + input_rows + half_rotary + rotary_offsets, mask=rotary_tile_mask, other=0.0 + ).to(tl.float32) + weight_first = tl.load( + weight_ptr + rotary_offsets, mask=rotary_offsets < half_rotary, other=0.0 + ).to(tl.float32) + weight_second = tl.load( + weight_ptr + half_rotary + rotary_offsets, + mask=rotary_offsets < half_rotary, + other=0.0, + ).to(tl.float32) + x_first = (x_first * inverse_rms * (weight_first + 1.0)).to(output_dtype).to(tl.float32) + x_second = ( + (x_second * inverse_rms * (weight_second + 1.0)).to(output_dtype).to(tl.float32) + ) + else: + # Full power-of-two RoPE can split normalized registers into NeoX halves directly, + # avoiding another set of input and weight loads. + x_first, x_second = tl.split( + tl.permute(tl.reshape(normalized, (block_tokens, 2, half_rotary)), (0, 2, 1)) + ) if use_mrope: + # Interleave height/width frequencies within their section lengths; remaining + # frequencies use temporal positions (section zero). section = tl.where( (rotary_offsets % 3 == 1) & (rotary_offsets < mrope_section1 * 3), 1, - tl.where( - (rotary_offsets % 3 == 2) & (rotary_offsets < mrope_section2 * 3), - 2, - 0, - ), + tl.where((rotary_offsets % 3 == 2) & (rotary_offsets < mrope_section2 * 3), 2, 0), ) position = tl.load( - positions_ptr + section * num_tokens + token, - mask=rotary_mask, + positions_ptr + section * num_tokens + tokens, + mask=rotary_tile_mask, other=0, ).to(tl.int64) else: - position = tl.load(positions_ptr + token).to(tl.int64) + position = tl.load(positions_ptr + tokens, mask=token_mask, other=0).to(tl.int64) position_is_valid = (position >= 0) & (position < max_positions) - tl.device_assert(position_is_valid, "position is outside the RoPE table") + tl.device_assert( + position_is_valid | ~rotary_tile_mask, "position is outside the RoPE table" + ) + # Assertions may be disabled: keep table loads safe, then propagate NaNs + # for invalid positions rather than silently using position zero. safe_position = tl.where(position_is_valid, position, 0) cos_sin_base = cos_sin_ptr + safe_position * cos_sin_stride_position - cos = tl.load(cos_sin_base + rotary_offsets, mask=rotary_mask, other=0.0).to(tl.float32) - sin = tl.load(cos_sin_base + half_rotary + rotary_offsets, mask=rotary_mask, other=0.0).to( + cos = tl.load(cos_sin_base + rotary_offsets, mask=rotary_tile_mask, other=0.0).to( tl.float32 ) + sin = tl.load( + cos_sin_base + half_rotary + rotary_offsets, mask=rotary_tile_mask, other=0.0 + ).to(tl.float32) cos = tl.where(position_is_valid, cos, float("nan")) sin = tl.where(position_is_valid, sin, float("nan")) tl.store( - output_base + rotary_offsets, + output_rows + rotary_offsets, x_first * cos - x_second * sin, - mask=rotary_mask, + mask=rotary_tile_mask, ) tl.store( - output_base + half_rotary + rotary_offsets, + output_rows + half_rotary + rotary_offsets, x_second * cos + x_first * sin, - mask=rotary_mask, + mask=rotary_tile_mask, ) if not is_k: - gate_input = input_base + head_dim - gate_output = gate_out_ptr + token * gate_out_stride_token + local_head * head_dim - gate = tl.load(gate_input + head_offsets, mask=head_mask, other=0.0) - tl.store(gate_output + head_offsets, gate, mask=head_mask) + # Gates bypass normalization and RoPE; preserve their input bits. + gate = tl.load(input_rows + head_dim + head_offsets, mask=tile_mask, other=0.0) + gate_rows = gate_out_ptr + tokens * gate_out_stride_token + local_head * head_dim + tl.store(gate_rows + head_offsets, gate, mask=tile_mask) else: + # V heads only move from the interleaved projection to packed QKV. local_head = head - num_q_heads - num_kv_heads - head_offsets = tl.arange(0, head_block) - head_mask = head_offsets < head_dim - input_base = ( - qkv_ptr + token * qkv_stride_token + 2 * q_size + kv_size + local_head * head_dim + value = tl.load( + qkv_ptr + + tokens * qkv_stride_token + + 2 * q_size + + kv_size + + local_head * head_dim + + head_offsets, + mask=tile_mask, + other=0.0, ) - output_base = ( - qkv_out_ptr + token * qkv_out_stride_token + q_size + kv_size + local_head * head_dim + tl.store( + qkv_out_ptr + + tokens * qkv_out_stride_token + + q_size + + kv_size + + local_head * head_dim + + head_offsets, + value, + mask=tile_mask, ) - value = tl.load(input_base + head_offsets, mask=head_mask, other=0.0) - tl.store(output_base + head_offsets, value, mask=head_mask) + + +def _use_tiled_qkv_rope( + qkv: torch.Tensor, + num_q_heads: int, + num_kv_heads: int, + head_dim: int, + rotary_dim: int, + enable_tiling: bool, +) -> bool: + return ( + # Opt in because changing tile layouts can change Q/K rounding. + enable_tiling + # Restrict tiling to the full- and partial-RoPE geometries benchmarked. + and (head_dim, rotary_dim) in ((128, 128), (256, 64)) + # Count Q, K, and V head rows so small TP-local head counts still need + # enough work: 3072 rows yield at least 384 eight-token programs. + and qkv.shape[0] * (num_q_heads + 2 * num_kv_heads) >= _MIN_TILED_HEAD_ROWS + # Avoid measured small-batch regressions. Larger inputs only add tiles. + and qkv.shape[0] >= 256 + ) @torch.library.custom_op( @@ -180,6 +259,7 @@ def _fused_qkv_gemma_rmsnorm_rope_gate( use_mrope: bool, mrope_section1: int, mrope_section2: int, + enable_tiling: bool, ) -> Tuple[torch.Tensor, torch.Tensor]: assert qkv.dim() == 2 and qkv.stride(-1) == 1 assert qkv.dtype in (torch.bfloat16, torch.float16) @@ -214,7 +294,16 @@ def _fused_qkv_gemma_rmsnorm_rope_gate( half_rotary = rotary_dim // 2 head_block = triton.next_power_of_2(head_dim) rotary_block = triton.next_power_of_2(half_rotary) - grid = (num_tokens, num_q_heads + 2 * num_kv_heads) + # Eight tokens per program amortize scheduling overhead across independent rows. + # Two warps are the benchmarked configuration for eligible workloads; the scalar fallback keeps + # one token and four warps. + block_tokens, num_warps = ( + (8, 2) + if _use_tiled_qkv_rope(qkv, num_q_heads, num_kv_heads, head_dim, rotary_dim, enable_tiling) + else (1, 4) + ) + # Eight-token tiles reduce program count; token masks handle the final tile. + grid = (triton.cdiv(num_tokens, block_tokens), num_q_heads + 2 * num_kv_heads) _fused_qkv_gemma_rmsnorm_rope_gate_kernel[grid]( qkv, qkv_out, @@ -242,7 +331,8 @@ def _fused_qkv_gemma_rmsnorm_rope_gate( use_mrope=use_mrope, mrope_section1=mrope_section1, mrope_section2=mrope_section2, - num_warps=4, + block_tokens=block_tokens, + num_warps=num_warps, ) return qkv_out, gate_out @@ -262,6 +352,7 @@ def _( use_mrope: bool, mrope_section1: int, mrope_section2: int, + enable_tiling: bool, ) -> Tuple[torch.Tensor, torch.Tensor]: num_tokens = qkv.shape[0] qkv_width = (num_q_heads + 2 * num_kv_heads) * head_dim @@ -283,9 +374,23 @@ def fused_qkv_gemma_rmsnorm_rope_gate( head_dim: int, rotary_dim: int, mrope_section: Optional[Tuple[int, int, int]] = None, + *, + enable_tiling: bool = _QKV_ROPE_TILING_ENABLED, ) -> Tuple[torch.Tensor, torch.Tensor]: """Prepare packed QKV and gate with one Triton kernel. + `TRTLLM_ENABLE_QKV_ROPE_TILING=1` opts into eight-token tiles for at least 256 tokens with + (head dim, rotary dim) equal to (128, 128) or (256, 64) and at least 3072 token/head pairs + across Q/K/V. + Head counts can vary with tensor parallelism. The flag sets the default at import; callers can + override it with `enable_tiling`. + Other calls use the single-token kernel. + + Tiled FP32 reductions and compiler contraction can change Q/K rounding. With tiling enabled, + changing a request's batch size across a dispatch boundary can change its Q/K values, cached K, + and generated tokens. + Leave the flag off when batch-invariant results are required. + Args: qkv: ``[num_tokens, 2 * q_size + 2 * kv_size]`` BF16/FP16 projection output in per-head interleaved Q/G layout. @@ -301,6 +406,8 @@ def fused_qkv_gemma_rmsnorm_rope_gate( rotary_dim: Prefix dimension receiving NeoX RoPE. mrope_section: Temporal/height/width rotary-half dimensions. ``None`` selects plain RoPE; a tuple selects Qwen-style interleaved MRoPE. + enable_tiling: Allow eight-token tiling for eligible workloads. Defaults + to the import-time environment flag; ``True`` does not bypass eligibility. Returns: Packed QKV ``[num_tokens, q_size + 2 * kv_size]`` and gate @@ -329,6 +436,7 @@ def fused_qkv_gemma_rmsnorm_rope_gate( use_mrope, mrope_section1, mrope_section2, + enable_tiling, ) diff --git a/tests/unittest/_torch/attention/kernels/test_fused_qk_norm_rope_gate.py b/tests/unittest/_torch/attention/kernels/test_fused_qk_norm_rope_gate.py index 414bb7868d76..3c0a5946238a 100644 --- a/tests/unittest/_torch/attention/kernels/test_fused_qk_norm_rope_gate.py +++ b/tests/unittest/_torch/attention/kernels/test_fused_qk_norm_rope_gate.py @@ -13,9 +13,14 @@ # See the License for the specific language governing permissions and # limitations under the License. +import json +from pathlib import Path +from unittest.mock import patch + import pytest import torch +import tensorrt_llm._torch.attention.kernels.fused_qk_norm_rope_gate as _kernel_module from tensorrt_llm._torch.attention.attention import Attention from tensorrt_llm._torch.attention.kernels.fused_qk_norm_rope_gate import ( fused_qkv_gemma_rmsnorm_rope_gate, @@ -23,6 +28,12 @@ ) +@pytest.fixture(params=(False, True), ids=("stock", "tiled")) +def qkv_launch(request: pytest.FixtureRequest, monkeypatch: pytest.MonkeyPatch) -> None: + # Exercise both kernels on every GPU, independently of production eligibility. + monkeypatch.setattr(_kernel_module, "_use_tiled_qkv_rope", lambda *args: request.param) + + def _make_cos_sin(max_positions: int, rotary_dim: int, theta: float = 1_000_000.0): inverse_frequency = 1.0 / ( theta ** (torch.arange(0, rotary_dim, 2, dtype=torch.float32, device="cuda") / rotary_dim) @@ -99,6 +110,7 @@ def _rope(x): (333, 16, 4, 256, 256), ), ) +@pytest.mark.usefixtures("qkv_launch") def test_fused_qkv_gemma_rmsnorm_rope_gate_matches_reference( dtype, num_tokens, num_q_heads, num_kv_heads, head_dim, rotary_dim ): @@ -159,6 +171,7 @@ def test_fused_sigmoid_mul_inplace_supports_strided_gate(): @pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) +@pytest.mark.usefixtures("qkv_launch") def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_interleaved_mrope(dtype): torch.manual_seed(2026) num_tokens, num_q_heads, num_kv_heads = 37, 16, 2 @@ -208,6 +221,7 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_interleaved_mrope(dtype): @pytest.mark.parametrize("num_tokens", (1, 37, 333)) +@pytest.mark.usefixtures("qkv_launch") def test_fused_qkv_gemma_rmsnorm_rope_gate_matches_production_thop(num_tokens): torch.manual_seed(9027) num_q_heads, num_kv_heads = 16, 2 @@ -295,9 +309,13 @@ def test_fused_sigmoid_mul_inplace_supports_flat_gate(): torch.testing.assert_close(actual.float(), expected, atol=0.005, rtol=0.005) -def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(): +@pytest.mark.parametrize("enabled", (False, True)) +@pytest.mark.parametrize("backend", ("eager", "inductor")) +def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile( + backend: str, enabled: bool +) -> None: torch.manual_seed(2027) - num_tokens, num_q_heads, num_kv_heads = 7, 8, 2 + num_tokens, num_q_heads, num_kv_heads = 257, 8, 2 head_dim, rotary_dim = 128, 128 q_size = num_q_heads * head_dim kv_size = num_kv_heads * head_dim @@ -308,12 +326,12 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(): ) q_weight = torch.randn((head_dim,), dtype=torch.bfloat16, device="cuda") * 0.1 k_weight = torch.randn((head_dim,), dtype=torch.bfloat16, device="cuda") * 0.1 - cos_sin = _make_cos_sin(128, rotary_dim) + cos_sin = _make_cos_sin(num_tokens, rotary_dim) positions = torch.arange(num_tokens, dtype=torch.int32, device="cuda") compiled = torch.compile( fused_qkv_gemma_rmsnorm_rope_gate, - backend="eager", + backend=backend, fullgraph=True, ) actual = compiled( @@ -327,6 +345,7 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(): num_kv_heads, head_dim, rotary_dim, + enable_tiling=enabled, ) expected = fused_qkv_gemma_rmsnorm_rope_gate( qkv, @@ -339,6 +358,7 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(): num_kv_heads, head_dim, rotary_dim, + enable_tiling=enabled, ) torch.testing.assert_close(actual, expected, atol=0, rtol=0) @@ -356,3 +376,210 @@ def apply_gate(attention_output, gate_input): compiled = torch.compile(apply_gate, backend="eager", fullgraph=True) actual = compiled(attention.clone(), gate) torch.testing.assert_close(actual, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) +@pytest.mark.parametrize("num_tokens", (0, 1, 7, 8, 9, 255, 256, 257, 513)) +@pytest.mark.parametrize( + "q_heads,kv_heads,head_dim,rotary_dim", + ((8, 2, 128, 128), (16, 2, 256, 64), (8, 1, 256, 64), (8, 2, 192, 96), (8, 2, 192, 192)), +) +def test_tiled_qkv_rope_matches_stock_and_replays( + dtype: torch.dtype, num_tokens: int, q_heads: int, kv_heads: int, head_dim: int, rotary_dim: int +) -> None: + torch.manual_seed(2028) + width = 2 * (q_heads + kv_heads) * head_dim + qkv = torch.randn((num_tokens, width + 37), dtype=dtype, device="cuda")[:, :width] + qw = torch.randn(head_dim, dtype=dtype, device="cuda") * 0.1 + kw = torch.randn_like(qw) * 0.1 + table = _make_cos_sin(128, rotary_dim) + positions = torch.randint(0, 128, (3, num_tokens), dtype=torch.int64, device="cuda") + half = rotary_dim // 2 + sections = (half - 2 * (half // 3), half // 3, half // 3) + + def run() -> tuple[torch.Tensor, torch.Tensor]: + return fused_qkv_gemma_rmsnorm_rope_gate( + qkv, + qw, + kw, + table, + positions, + 1e-6, + q_heads, + kv_heads, + head_dim, + rotary_dim, + sections, + ) + + with patch.object(_kernel_module, "_use_tiled_qkv_rope", return_value=False): + expected = run() + with patch.object(_kernel_module, "_use_tiled_qkv_rope", return_value=True): + actual = run() + torch.testing.assert_close(actual[0], expected[0], atol=0.02, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=0, rtol=0) + v_start = (q_heads + kv_heads) * head_dim + torch.testing.assert_close(actual[0][:, v_start:], expected[0][:, v_start:], atol=0, rtol=0) + if num_tokens == 0: + return + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + for _ in range(3): + qkv.add_(0.125) + expected = run() + graph.replay() + torch.testing.assert_close(captured, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize("enabled", (False, True)) +def test_qkv_rope_tiling_dispatch_is_architecture_independent( + enabled: bool, +) -> None: + qkv = torch.empty((256, 1024), device="cuda") + with patch.object(torch.cuda, "get_device_properties") as query: + assert _kernel_module._use_tiled_qkv_rope(qkv, 16, 2, 256, 64, enabled) == enabled + query.assert_not_called() + + +@pytest.mark.parametrize( + "num_tokens,head_dim,rotary_dim", + ((0, 256, 64), (255, 256, 64), (256, 128, 64), (256, 192, 96), (256, 256, 256)), +) +def test_qkv_rope_tiling_excluded_shapes(num_tokens: int, head_dim: int, rotary_dim: int) -> None: + qkv = torch.empty((num_tokens, 1024), device="cuda") + assert not _kernel_module._use_tiled_qkv_rope(qkv, 16, 2, head_dim, rotary_dim, True) + + +@pytest.mark.parametrize("enabled", (False, True)) +@pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) +@pytest.mark.parametrize("q_heads,dim,rotary", ((16, 256, 64), (8, 128, 128))) +@pytest.mark.parametrize("tokens", (257, 32737, 32768, 32769)) +def test_qkv_rope_tiling_production_dispatch( + tmp_path: Path, + enabled: bool, + dtype: torch.dtype, + q_heads: int, + dim: int, + rotary: int, + tokens: int, +) -> None: + kv_heads = 2 + qkv = torch.randn((tokens, 2 * (q_heads + kv_heads) * dim), device="cuda", dtype=dtype) + weight = torch.randn(dim, device="cuda", dtype=dtype) * 0.1 + table = _make_cos_sin(tokens, rotary) + positions = torch.arange(tokens, device="cuda", dtype=torch.int32) + + def run() -> tuple[torch.Tensor, torch.Tensor]: + return fused_qkv_gemma_rmsnorm_rope_gate( + qkv, + weight, + weight, + table, + positions, + 1e-6, + q_heads, + kv_heads, + dim, + rotary, + enable_tiling=enabled, + ) + + run() + torch.cuda.synchronize() + with torch.profiler.profile(activities=[torch.profiler.ProfilerActivity.CUDA]) as profile: + actual = run() + torch.cuda.synchronize() + trace_file = tmp_path / "dispatch.json" + profile.export_chrome_trace(str(trace_file)) + launches = [ + event + for event in json.loads(trace_file.read_text())["traceEvents"] + if event.get("cat") == "kernel" + and "_fused_qkv_gemma_rmsnorm_rope_gate_kernel" in event["name"] + ] + block_tokens = 8 if enabled else 1 + assert len(launches) == 1 + assert launches[0]["args"]["grid"] == [ + (tokens + block_tokens - 1) // block_tokens, + q_heads + 2 * kv_heads, + 1, + ] + expected = _reference( + qkv, weight, weight, table, positions, 1e-6, q_heads, kv_heads, dim, rotary + ) + torch.testing.assert_close(actual[0], expected[0], atol=0.02, rtol=0.02) + torch.testing.assert_close(actual[1], expected[1], atol=0, rtol=0) + v_start = (q_heads + kv_heads) * dim + torch.testing.assert_close(actual[0][:, v_start:], expected[0][:, v_start:], atol=0, rtol=0) + if tokens == 32737: + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + captured = run() + qkv.add_(0.125) + expected = run() + graph.replay() + torch.testing.assert_close(captured, expected, atol=0, rtol=0) + + +@pytest.mark.parametrize( + "tokens,q_heads,kv_heads,expected", + ( + (256, 16, 2, True), + (256, 8, 1, False), + (512, 8, 1, True), + (256, 32, 4, True), + (512, 4, 1, True), + (1024, 1, 1, True), + ), +) +def test_qkv_rope_tiling_accepts_tp_head_counts( + tokens: int, q_heads: int, kv_heads: int, expected: bool +) -> None: + qkv = torch.empty((tokens, 2 * (q_heads + kv_heads) * 256), device="cuda", dtype=torch.bfloat16) + assert _kernel_module._use_tiled_qkv_rope(qkv, q_heads, kv_heads, 256, 64, True) == expected + + +@pytest.mark.parametrize("dtype", (torch.bfloat16, torch.float16)) +@pytest.mark.parametrize("q_heads,head_dim,rotary_dim", ((8, 128, 128), (16, 256, 64))) +def test_qkv_rope_tiling_batch_boundaries( + dtype: torch.dtype, + q_heads: int, + head_dim: int, + rotary_dim: int, +) -> None: + torch.manual_seed(2030) + kv_heads, prefix_tokens = 2, 255 + qkv = torch.randn((32769, 2 * (q_heads + kv_heads) * head_dim), device="cuda", dtype=dtype) + weight = torch.randn(head_dim, device="cuda", dtype=dtype) * 0.1 + table = _make_cos_sin(128, rotary_dim) + positions = torch.randint(0, 128, (3, 32769), device="cuda", dtype=torch.int32) + half = rotary_dim // 2 + sections = (half - 2 * (half // 3), half // 3, half // 3) + reference = None + for tokens in (255, 256, 257, 16384, 16385, 32737, 32768, 32769): + actual = fused_qkv_gemma_rmsnorm_rope_gate( + qkv[:tokens], + weight, + weight, + table, + positions[:, :tokens].contiguous(), + 1e-6, + q_heads, + kv_heads, + head_dim, + rotary_dim, + sections, + enable_tiling=True, + ) + prefix = tuple(value[:prefix_tokens].clone() for value in actual) + if reference is None: + reference = prefix + else: + # The opt-in permits Q/K rounding changes at dispatch boundaries. + torch.testing.assert_close(prefix[0], reference[0], atol=0.02, rtol=0.02) + torch.testing.assert_close(prefix[1], reference[1], atol=0, rtol=0) + v_start = (q_heads + kv_heads) * head_dim + torch.testing.assert_close( + prefix[0][:, v_start:], reference[0][:, v_start:], atol=0, rtol=0 + )