[None][perf] Add opt-in tiled QKV preprocessing

docs/source/deployment-guide/deployment-guide-for-qwen3.8-qwen3.5-on-trtllm.md

@@ -1,3 +1,8 @@
+<!--
+SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+SPDX-License-Identifier: Apache-2.0
+-->
+
 # 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.

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,
     )