From 1cccf20ac000bbfc4875e7de62040e5f6ef8bd4a Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Tue, 27 Oct 2026 23:07:29 -0700 Subject: [PATCH 1/2] [None][perf] Add opt-in tiled QKV preprocessing on Rubin * Why? One CTA per token and head adds overhead to large Qwen attention batches. Token tiling reduces this overhead while retaining the fused normalization, rotary embedding, and gate extraction. * What? Add an eight-token path behind TRTLLM_ENABLE_QKV_ROPE_TILING=1. Limit selection to validated shapes with at least 256 tokens. Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- .../kernels/fused_qk_norm_rope_gate.py | 224 ++++++++++++----- .../kernels/test_fused_qk_norm_rope_gate.py | 225 +++++++++++++++++- 2 files changed, 384 insertions(+), 65 deletions(-) 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..df3d81dbcc5a 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,178 @@ 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 +) -> bool: + return ( + # Opt in because changing tile layouts can change Q/K rounding. + _QKV_ROPE_TILING_ENABLED + # 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( @@ -214,7 +288,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) + 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 +325,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 @@ -286,6 +370,20 @@ def fused_qkv_gemma_rmsnorm_rope_gate( ) -> 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 is read at import. + Other calls use the single-token kernel. + + Tiling is benchmarked on B200 and Rubin. Other GPU architectures can opt in, + but their performance is not established by those measurements. + + 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. 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..9ddf613224e2 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,7 +309,9 @@ 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.usefixtures("qkv_launch") +@pytest.mark.parametrize("backend", ("eager", "inductor")) +def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(backend: str) -> None: torch.manual_seed(2027) num_tokens, num_q_heads, num_kv_heads = 7, 8, 2 head_dim, rotary_dim = 128, 128 @@ -313,7 +329,7 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(): compiled = torch.compile( fused_qkv_gemma_rmsnorm_rope_gate, - backend="eager", + backend=backend, fullgraph=True, ) actual = compiled( @@ -356,3 +372,208 @@ 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( + monkeypatch: pytest.MonkeyPatch, enabled: bool +) -> None: + monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", enabled) + 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 + 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( + monkeypatch: pytest.MonkeyPatch, num_tokens: int, head_dim: int, rotary_dim: int +) -> None: + monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", True) + qkv = torch.empty((num_tokens, 1024), device="cuda") + assert not _kernel_module._use_tiled_qkv_rope(qkv, 16, 2, head_dim, rotary_dim) + + +@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( + monkeypatch: pytest.MonkeyPatch, + tmp_path: Path, + enabled: bool, + dtype: torch.dtype, + q_heads: int, + dim: int, + rotary: int, + tokens: int, +) -> None: + monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", enabled) + 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 + ) + + 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( + monkeypatch: pytest.MonkeyPatch, tokens: int, q_heads: int, kv_heads: int, expected: bool +) -> None: + monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", True) + 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) == 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( + monkeypatch: pytest.MonkeyPatch, + dtype: torch.dtype, + q_heads: int, + head_dim: int, + rotary_dim: int, +) -> None: + monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", True) + 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, + ) + 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 + ) From a50e4bf0fc6b7f895ab2ef27e03ff1dc928d4be7 Mon Sep 17 00:00:00 2001 From: William Zhang <133824995+2ez4bz@users.noreply.github.com> Date: Wed, 28 Oct 2026 14:00:22 -0700 Subject: [PATCH 2/2] address review comments Signed-off-by: William Zhang <133824995+2ez4bz@users.noreply.github.com> --- ...ent-guide-for-qwen3.8-qwen3.5-on-trtllm.md | 13 ++++++ .../kernels/fused_qk_norm_rope_gate.py | 36 +++++++++------ .../kernels/test_fused_qk_norm_rope_gate.py | 46 +++++++++++-------- 3 files changed, 62 insertions(+), 33 deletions(-) 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 df3d81dbcc5a..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 @@ -221,11 +221,16 @@ def _fused_qkv_gemma_rmsnorm_rope_gate_kernel( def _use_tiled_qkv_rope( - qkv: torch.Tensor, num_q_heads: int, num_kv_heads: int, head_dim: int, rotary_dim: int + 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. - _QKV_ROPE_TILING_ENABLED + 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 @@ -254,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) @@ -293,7 +299,7 @@ def _fused_qkv_gemma_rmsnorm_rope_gate( # 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) + 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. @@ -346,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 @@ -367,21 +374,21 @@ 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 is read at import. + `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. - Tiling is benchmarked on B200 and Rubin. Other GPU architectures can opt in, - but their performance is not established by those measurements. - - 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. + 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: @@ -399,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 @@ -427,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 9ddf613224e2..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 @@ -309,11 +309,13 @@ def test_fused_sigmoid_mul_inplace_supports_flat_gate(): torch.testing.assert_close(actual.float(), expected, atol=0.005, rtol=0.005) -@pytest.mark.usefixtures("qkv_launch") +@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) -> None: +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 @@ -324,7 +326,7 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(backend: str) ) 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( @@ -343,6 +345,7 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(backend: str) num_kv_heads, head_dim, rotary_dim, + enable_tiling=enabled, ) expected = fused_qkv_gemma_rmsnorm_rope_gate( qkv, @@ -355,6 +358,7 @@ def test_fused_qkv_gemma_rmsnorm_rope_gate_supports_torch_compile(backend: str) num_kv_heads, head_dim, rotary_dim, + enable_tiling=enabled, ) torch.testing.assert_close(actual, expected, atol=0, rtol=0) @@ -430,12 +434,11 @@ def run() -> tuple[torch.Tensor, torch.Tensor]: @pytest.mark.parametrize("enabled", (False, True)) def test_qkv_rope_tiling_dispatch_is_architecture_independent( - monkeypatch: pytest.MonkeyPatch, enabled: bool + enabled: bool, ) -> None: - monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", enabled) 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 + assert _kernel_module._use_tiled_qkv_rope(qkv, 16, 2, 256, 64, enabled) == enabled query.assert_not_called() @@ -443,12 +446,9 @@ def test_qkv_rope_tiling_dispatch_is_architecture_independent( "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( - monkeypatch: pytest.MonkeyPatch, num_tokens: int, head_dim: int, rotary_dim: int -) -> None: - monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", True) +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) + assert not _kernel_module._use_tiled_qkv_rope(qkv, 16, 2, head_dim, rotary_dim, True) @pytest.mark.parametrize("enabled", (False, True)) @@ -456,7 +456,6 @@ def test_qkv_rope_tiling_excluded_shapes( @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( - monkeypatch: pytest.MonkeyPatch, tmp_path: Path, enabled: bool, dtype: torch.dtype, @@ -465,7 +464,6 @@ def test_qkv_rope_tiling_production_dispatch( rotary: int, tokens: int, ) -> None: - monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", enabled) 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 @@ -474,7 +472,17 @@ def test_qkv_rope_tiling_production_dispatch( 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 + qkv, + weight, + weight, + table, + positions, + 1e-6, + q_heads, + kv_heads, + dim, + rotary, + enable_tiling=enabled, ) run() @@ -526,23 +534,20 @@ def run() -> tuple[torch.Tensor, torch.Tensor]: ), ) def test_qkv_rope_tiling_accepts_tp_head_counts( - monkeypatch: pytest.MonkeyPatch, tokens: int, q_heads: int, kv_heads: int, expected: bool + tokens: int, q_heads: int, kv_heads: int, expected: bool ) -> None: - monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", True) 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) == expected + 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( - monkeypatch: pytest.MonkeyPatch, dtype: torch.dtype, q_heads: int, head_dim: int, rotary_dim: int, ) -> None: - monkeypatch.setattr(_kernel_module, "_QKV_ROPE_TILING_ENABLED", True) 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) @@ -565,6 +570,7 @@ def test_qkv_rope_tiling_batch_boundaries( head_dim, rotary_dim, sections, + enable_tiling=True, ) prefix = tuple(value[:prefix_tokens].clone() for value in actual) if reference is None: