From 2c75d417a525da936ffb792383d21537643b48c9 Mon Sep 17 00:00:00 2001 From: Dan Sun Date: Sat, 15 Aug 2026 22:07:28 -0400 Subject: [PATCH] [None][feat] Run Gemma4 head_dim 512 attention on Triton off Blackwell Gemma4's full-attention layers use head_dim=512, which no SM90 paged attention kernel serves: trtllm-gen has the cubins but ships them for datacenter Blackwell only, TRTLLM paged FMHA does not cover 512, and FlashInfer's fa2/fa3 paged kernels stop at 256. Gemma4 therefore does not run on Hopper through the PyTorch backend at all. Route those layers through Triton instead. The context phase already had a suitable kernel in triton_prefill.py (head_dim 512 capable, with a Hopper tile heuristic); the decode half is ported from the AutoDeploy Triton attention backend, which uses the identical combined HND cache layout [num_pages, 2, num_kv_heads, page_size, head_dim]. Both phases are handled before metadata.plan(), so no FlashInfer wrapper is ever created for these layers and they never touch workspace_buffer. That keeps a single wrapper type in play across the rest of the model and avoids the workspace corruption that mixing wrapper types under CUDA graphs causes. Sliding layers (head_dim 256) are unaffected and stay on FlashInfer. KV-shared layers, speculative-decoding draft views and multi-token generation steps raise NotImplementedError on this path rather than returning incorrect results. CUDA graph capture and perf tuning are left to a follow-up. Signed-off-by: Dan Sun --- .../_torch/attention_backend/flashinfer.py | 110 +++- .../_torch/attention_backend/triton_decode.py | 488 ++++++++++++++++++ .../test_lists/test-db/l0_h100.yml | 9 + .../_torch/attention/backend_capability.py | 21 +- .../_torch/attention/test_triton_decode.py | 388 ++++++++++++++ .../_torch/modeling/test_modeling_gemma4.py | 111 ++++ 6 files changed, 1125 insertions(+), 2 deletions(-) create mode 100644 tensorrt_llm/_torch/attention_backend/triton_decode.py create mode 100644 tests/unittest/_torch/attention/test_triton_decode.py diff --git a/tensorrt_llm/_torch/attention_backend/flashinfer.py b/tensorrt_llm/_torch/attention_backend/flashinfer.py index d49ab084a128..a19df8256df5 100644 --- a/tensorrt_llm/_torch/attention_backend/flashinfer.py +++ b/tensorrt_llm/_torch/attention_backend/flashinfer.py @@ -33,7 +33,7 @@ from flashinfer.jit.core import check_cuda_arch from typing_extensions import Self -from tensorrt_llm._utils import nvtx_range +from tensorrt_llm._utils import is_sm_100f, nvtx_range from tensorrt_llm.functional import AttentionMaskType from tensorrt_llm.logger import logger from tensorrt_llm.models.modeling_utils import QuantConfig @@ -64,6 +64,14 @@ _MAX_CUDA_THREADS_PER_BLOCK = 1024 +_FLASHINFER_MAX_PAGED_HEAD_DIM = 256 +"""Largest head_dim FlashInfer's fa2/fa3 paged kernels handle. + +Beyond this only trtllm-gen has cubins, and those ship for datacenter Blackwell +only. Layers above this threshold run through Triton on other architectures -- +see the ``use_triton_attention`` path in ``FlashInferAttention.forward_impl``. +""" + def _slice_paged_kv_cache_heads( paged_kv_cache: torch.Tensor | tuple[torch.Tensor, torch.Tensor], @@ -2335,6 +2343,106 @@ def decode_forward(plan_params: PlanParams, out: torch.Tensor): kv_cache, out=out.view(-1, self.num_heads, self.head_dim)) + # Triton attention for layers FlashInfer's paged kernels cannot serve + # at all. fa2/fa3 cover head_dim<=256; beyond that only trtllm-gen has + # cubins, and those are datacenter-Blackwell only. On every other + # architecture (Hopper included) run both phases through Triton. + # + # Both phases are handled here so that no FlashInfer wrapper is ever + # created for these layers: the Triton kernels use their own scratch + # and never touch ``workspace_buffer``. That keeps a single wrapper + # type in play across the rest of the model, avoiding the workspace + # corruption that mixing wrapper types under CUDA graphs causes. + use_triton_attention = (self.head_dim > _FLASHINFER_MAX_PAGED_HEAD_DIM + and not is_sm_100f()) + + if use_triton_attention: + if (metadata._is_shared_kv_draft_view + or metadata._is_separate_kv_draft_view): + raise NotImplementedError( + "Speculative decoding draft metadata views require the " + "trtllm-gen decode backend, which has no cubins for " + f"head_dim {self.head_dim} on this architecture.") + if num_contexts > 0 and k is None: + # KV-shared layers read another layer's cache and pass k=None, + # so the current tokens' KV is already paged. Feeding them as + # Triton's "prefix" would drop causal masking between them, so + # refuse rather than return wrong numbers. + raise NotImplementedError( + "KV-shared layers are not yet supported on the Triton " + f"attention path (head_dim {self.head_dim}, layer " + f"{self.layer_idx}). Use a model without " + "num_kv_shared_layers, or run on datacenter Blackwell.") + + from .triton_decode import triton_decode + from .triton_prefill import triton_prefill_with_custom_mask + + logger.info_once( + "FlashInfer paged kernels do not cover head_dim " + f"{self.head_dim}; using the Triton prefill/decode path.", + key=f"triton_attention_hd{self.head_dim}") + + sm_scale = 1 / (math.sqrt(self.head_dim) * self.q_scaling) + # attention_window_size is already in flashinfer convention here + # (forward() subtracts 1 from the exclusive TRTLLM window). + window_left = (attention_window_size + if attention_window_size is not None else -1) + + if num_contexts > 0: + triton_prefill_with_custom_mask( + q=q[:num_ctx_tokens], + k=k[:num_ctx_tokens], + v=v[:num_ctx_tokens], + output=output[:num_ctx_tokens].view(-1, self.num_heads, + self.head_dim), + qo_indptr=metadata.qo_indptr[:num_contexts + 1], + kv_cache=kv_cache, + prefix_lens=metadata. + cached_token_lens[:num_contexts].clone(), + page_table_indptr=metadata. + paged_kv_indptr_prefill[:num_contexts + 1], + page_table_indices=metadata. + _paged_kv_indices[:metadata.num_context_blocks], + page_size=metadata.page_size, + custom_mask=attention_mask_data, + sm_scale=sm_scale, + window_left=window_left, + ) + + if num_generations > 0: + # triton_decode derives its batch size from q's leading dim, so + # it assumes exactly one query token per generation request. + num_gen_tokens = q.shape[0] - num_ctx_tokens + if num_gen_tokens != num_generations: + raise NotImplementedError( + "Triton decode expects one query token per generation " + f"request, got {num_gen_tokens} tokens for " + f"{num_generations} requests. Multi-token generation " + "steps are not supported on this path.") + # triton_decode's sliding_window counts the tokens attended to + # *including* the current one; window_left excludes it. + sliding_window = (attention_window_size + + 1 if attention_window_size is not None else + None) + # Same decode-slice convention as the FlashInfer decode plan: + # indptr rebased to the generation range, indices offset past + # the context blocks. + triton_decode( + q=q[num_ctx_tokens:], + kv_cache=kv_cache, + kv_indices=metadata. + paged_kv_indices[metadata.num_context_blocks:], + kv_indptr=metadata. + paged_kv_indptr_decode[:num_generations + 1], + kv_last_page_len=metadata. + paged_kv_last_page_len[metadata.num_contexts:], + sm_scale=sm_scale, + sliding_window=sliding_window, + out=output[num_ctx_tokens:].view(-1, self.num_heads, + self.head_dim), + ) + return + # Triton prefill fallback: trtllm-gen cannot handle custom # (bidirectional) attention masks for head_dim>256 layers. Use a # Triton prefill kernel for those layers during multimodal prefill, diff --git a/tensorrt_llm/_torch/attention_backend/triton_decode.py b/tensorrt_llm/_torch/attention_backend/triton_decode.py new file mode 100644 index 000000000000..237244c37d54 --- /dev/null +++ b/tensorrt_llm/_torch/attention_backend/triton_decode.py @@ -0,0 +1,488 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Triton FlashDecoding for paged KV cache. + +Used as a fallback when the FlashInfer paged decode kernels cannot handle the +layer's head_dim, e.g. Gemma4's head_dim=512 full-attention layers on +architectures without trtllm-gen cubins (anything that is not datacenter +Blackwell). Companion to ``triton_prefill.py``, which covers the context phase +for the same layers. + +Ported from the AutoDeploy Triton attention backend +(``tensorrt_llm/_torch/auto_deploy/custom_ops/attention/triton_attention.py``), +which uses the identical combined HND cache layout +``[num_pages, 2, num_kv_heads, page_size, head_dim]``. Stripped of the +AutoDeploy ``AttentionDescriptor``/``AttentionRegistry`` wiring and the +context-phase kernels; only the decode path is kept. + +The KV cache is read with a cast to the query dtype, so an FP8 cache is +dequantized in-kernel and needs no separate conversion pass. +""" + +import math +from typing import Optional + +import torch +import triton +import triton.language as tl + +# Cache SM count to avoid repeated get_device_properties calls +_NUM_SMS: Optional[int] = None + +# Narrowest K that tl.dot accepts; page tiles are widened to at least this. +_MIN_TL_DOT_K = 16 + + +def _get_num_sms() -> int: + """Get the number of SMs on the current GPU (cached).""" + global _NUM_SMS + if _NUM_SMS is None: + _NUM_SMS = torch.cuda.get_device_properties(0).multi_processor_count + return _NUM_SMS + + +def _get_page_block(page_size: int) -> int: + """Return the page tile width used by Triton dot kernels.""" + return max(_MIN_TL_DOT_K, 1 << (page_size - 1).bit_length()) + + +def _get_num_splits(max_seq_len: int, batch_size: int, n_kv_heads: int, page_size: int) -> int: + """Compute optimal number of KV splits for FlashDecoding. + + With GQA batching, the grid is (batch, n_kv_heads, num_splits). + We want enough blocks to saturate the GPU. + """ + if max_seq_len <= 0: + return 1 + + num_sms = _get_num_sms() + existing_parallelism = batch_size * n_kv_heads + + # Already enough parallelism + if existing_parallelism >= num_sms * 2: + return 1 + + # Target ~4 waves of thread blocks + target_blocks = num_sms * 4 + num_splits = max(1, (target_blocks + existing_parallelism - 1) // existing_parallelism) + + # Cap splits so each block has at least 2 pages of work. With fewer pages, + # the per-block overhead (Q load, accumulator init, partial_o/lse store, + # plus stage2 reduction cost) dominates the useful compute (page-loop + # iterations). 2 pages is a conservative lower bound to keep the + # overhead-to-work ratio acceptable. + max_pages = max_seq_len // page_size + max_splits = max(1, max_pages // 2) + num_splits = min(num_splits, max_splits) + + # Round to next power of 2 for Triton compile caching + if num_splits > 1: + num_splits = 2 ** math.ceil(math.log2(num_splits)) + + return min(num_splits, 128) + + +@triton.autotune( + configs=[ + triton.Config({}, num_warps=2, num_stages=2), + triton.Config({}, num_warps=2, num_stages=3), + triton.Config({}, num_warps=4, num_stages=2), + triton.Config({}, num_warps=4, num_stages=3), + triton.Config({}, num_warps=8, num_stages=2), + triton.Config({}, num_warps=8, num_stages=3), + ], + key=["HEAD_DIM", "PAGE_SIZE", "PAGE_BLOCK", "HEAD_RATIO_PADDED", "SLIDING_WINDOW"], +) +@triton.jit +def _flash_decode_stage1_kernel( + # Query input + q_ptr, + # KV cache (combined) + kv_cache_ptr, + # Page table + kv_indices_ptr, + kv_indptr_ptr, + kv_last_page_len_ptr, + # Intermediate outputs + partial_o_ptr, + partial_lse_ptr, + # Q strides: [batch, n_heads, head_dim] + q_stride_batch: tl.constexpr, + q_stride_head: tl.constexpr, + # Partial output strides: [batch, n_heads, num_splits, head_dim] + po_stride_batch: tl.constexpr, + po_stride_head: tl.constexpr, + po_stride_split: tl.constexpr, + # Partial LSE strides: [batch, n_heads, num_splits] + plse_stride_batch: tl.constexpr, + plse_stride_head: tl.constexpr, + plse_stride_split: tl.constexpr, + # Cache strides: [num_blocks, 2, n_kv_heads, page_size, head_dim] + cache_stride_block: tl.constexpr, + cache_stride_kv: tl.constexpr, + cache_stride_head: tl.constexpr, + cache_stride_token: tl.constexpr, + # Constants + SM_SCALE: tl.constexpr, + N_HEADS: tl.constexpr, + N_KV_HEADS: tl.constexpr, + HEAD_DIM: tl.constexpr, + PAGE_SIZE: tl.constexpr, + PAGE_BLOCK: tl.constexpr, + HEAD_RATIO: tl.constexpr, + HEAD_RATIO_PADDED: tl.constexpr, + NUM_SPLITS: tl.constexpr, + SLIDING_WINDOW: tl.constexpr = 0, +): + """ + Key optimizations: + - Loads KV once for HEAD_RATIO Q heads + - Iterates by page for contiguous memory access + - Splits KV sequence across blocks for GPU utilization + """ + batch_id = tl.program_id(axis=0) + kv_head_id = tl.program_id(axis=1) + split_id = tl.program_id(axis=2) + + # Get sequence info from page table + kv_page_start = tl.load(kv_indptr_ptr + batch_id) + kv_page_end = tl.load(kv_indptr_ptr + batch_id + 1) + num_pages = kv_page_end - kv_page_start + last_page_len = tl.load(kv_last_page_len_ptr + batch_id) + + # Sliding window: restrict attention to pages within the window. + # Compute the total sequence length and the first valid KV position. + seq_len = (num_pages - 1) * PAGE_SIZE + last_page_len + if SLIDING_WINDOW > 0: + first_valid_pos = tl.maximum(0, seq_len - SLIDING_WINDOW) + first_window_page = first_valid_pos // PAGE_SIZE + else: + first_valid_pos = 0 + first_window_page = 0 + + # Only split over pages within the window + window_pages = num_pages - first_window_page + pages_per_split = (window_pages + NUM_SPLITS - 1) // NUM_SPLITS + page_split_start = first_window_page + split_id * pages_per_split + page_split_end = tl.minimum(page_split_start + pages_per_split, num_pages) + + dhead_offsets = tl.arange(0, HEAD_DIM) + # Use padded range for Triton power-of-2 requirement; mask out-of-bounds heads + head_local = tl.arange(0, HEAD_RATIO_PADDED) + head_ids = kv_head_id * HEAD_RATIO + head_local + head_mask = head_local < HEAD_RATIO + + # Handle inactive splits (beyond the sequence's pages) + if page_split_start >= num_pages: + # Store zeros + -inf LSE for valid HEAD_RATIO Q heads only + po_offsets = ( + batch_id * po_stride_batch + + head_ids[:, None] * po_stride_head + + split_id * po_stride_split + + dhead_offsets[None, :] + ) + tl.store( + partial_o_ptr + po_offsets, + tl.zeros([HEAD_RATIO_PADDED, HEAD_DIM], dtype=tl.float32), + mask=head_mask[:, None], + ) + plse_offsets = ( + batch_id * plse_stride_batch + + head_ids * plse_stride_head + + split_id * plse_stride_split + ) + tl.store( + partial_lse_ptr + plse_offsets, + tl.zeros([HEAD_RATIO_PADDED], dtype=tl.float32) + float("-inf"), + mask=head_mask, + ) + return + + # Load Q for HEAD_RATIO heads sharing this KV head: [HEAD_RATIO_PADDED, HEAD_DIM] + # Padded rows get zeros, producing zero attention scores (harmless, never stored) + q_offsets = ( + batch_id * q_stride_batch + head_ids[:, None] * q_stride_head + dhead_offsets[None, :] + ) + q_all = tl.load(q_ptr + q_offsets, mask=head_mask[:, None], other=0.0) + + acc = tl.zeros([HEAD_RATIO_PADDED, HEAD_DIM], dtype=tl.float32) + m_i = tl.zeros([HEAD_RATIO_PADDED], dtype=tl.float32) + float("-inf") + l_i = tl.zeros([HEAD_RATIO_PADDED], dtype=tl.float32) + + num_pages_this_split = page_split_end - page_split_start + for local_page_idx in range(num_pages_this_split): + page_idx = page_split_start + local_page_idx + physical_page = tl.load(kv_indices_ptr + kv_page_start + page_idx) + + # Determine valid tokens in this page + is_last_page_of_seq = page_idx == (num_pages - 1) + valid_tokens = tl.where(is_last_page_of_seq, last_page_len, PAGE_SIZE) + + page_offsets = tl.arange(0, PAGE_BLOCK) + page_mask = page_offsets < valid_tokens + + # Compute cache offset (use int64 to avoid overflow when + # physical_page * stride > 2^31) + cache_base = ( + physical_page.to(tl.int64) * cache_stride_block + + kv_head_id * cache_stride_head + + page_offsets[:, None] * cache_stride_token + + dhead_offsets[None, :] + ) + page_mask_2d = page_mask[:, None] + + k = tl.load(kv_cache_ptr + cache_base, mask=page_mask_2d, other=0.0).to( + q_all.dtype + ) # [PAGE_BLOCK, HEAD_DIM]; cast from fp8 if kv cache is fp8 + v = tl.load( + kv_cache_ptr + cache_base + cache_stride_kv, + mask=page_mask_2d, + other=0.0, + ).to(k.dtype) # [PAGE_BLOCK, HEAD_DIM]; cast from fp8 if kv cache is fp8 + + # [HEAD_RATIO_PADDED, HEAD_DIM] @ [HEAD_DIM, PAGE_BLOCK] + # -> [HEAD_RATIO_PADDED, PAGE_BLOCK] + attn = tl.dot(q_all, tl.trans(k)) * SM_SCALE + + # Combine validity mask with sliding window mask + if SLIDING_WINDOW > 0: + global_pos = page_idx * PAGE_SIZE + page_offsets + window_mask = global_pos >= first_valid_pos + attn = tl.where(page_mask[None, :] & window_mask[None, :], attn, float("-inf")) + else: + attn = tl.where(page_mask[None, :], attn, float("-inf")) + + # Online softmax update (vectorized over HEAD_RATIO_PADDED) + m_ij = tl.max(attn, axis=1) # [HEAD_RATIO_PADDED] + m_i_new = tl.maximum(m_i, m_ij) + alpha = tl.exp(m_i - m_i_new) + p = tl.exp(attn - m_i_new[:, None]) # [HEAD_RATIO_PADDED, PAGE_BLOCK] + + # [HEAD_RATIO_PADDED, PAGE_BLOCK] @ [PAGE_BLOCK, HEAD_DIM] + # -> [HEAD_RATIO_PADDED, HEAD_DIM] + acc = tl.dot(p.to(v.dtype), v, acc=acc * alpha[:, None]) + l_i = l_i * alpha + tl.sum(p, axis=1) + m_i = m_i_new + + # Finalize: normalize and compute LSE + l_i_safe = tl.where(l_i == 0.0, 1.0, l_i) + partial_o_val = acc / l_i_safe[:, None] # [HEAD_RATIO_PADDED, HEAD_DIM] + lse_val = m_i + tl.log(l_i_safe) # [HEAD_RATIO_PADDED] + + # Store results for valid HEAD_RATIO Q heads only (masked 2D store) + po_offsets = ( + batch_id * po_stride_batch + + head_ids[:, None] * po_stride_head + + split_id * po_stride_split + + dhead_offsets[None, :] + ) + tl.store(partial_o_ptr + po_offsets, partial_o_val, mask=head_mask[:, None]) + + plse_offsets = ( + batch_id * plse_stride_batch + head_ids * plse_stride_head + split_id * plse_stride_split + ) + tl.store(partial_lse_ptr + plse_offsets, lse_val, mask=head_mask) + + +@triton.jit +def _flash_decode_stage2_kernel( + # Partial results + partial_o_ptr, + partial_lse_ptr, + # Final output + o_ptr, + # Partial output strides: [batch, n_heads, num_splits, head_dim] + po_stride_batch: tl.constexpr, + po_stride_head: tl.constexpr, + po_stride_split: tl.constexpr, + # Partial LSE strides: [batch, n_heads, num_splits] + plse_stride_batch: tl.constexpr, + plse_stride_head: tl.constexpr, + plse_stride_split: tl.constexpr, + # Output strides: [batch, n_heads, head_dim] + o_stride_batch: tl.constexpr, + o_stride_head: tl.constexpr, + # Constants + HEAD_DIM: tl.constexpr, + NUM_SPLITS: tl.constexpr, +): + """ + Each program combines results from all splits for one (batch, head) pair. + """ + batch_id = tl.program_id(axis=0) + head_id = tl.program_id(axis=1) + + dhead_offsets = tl.arange(0, HEAD_DIM) + + # Find global maximum LSE across splits for numerical stability + global_max_lse = float("-inf") + for split_id in range(NUM_SPLITS): + plse_offset = ( + batch_id * plse_stride_batch + head_id * plse_stride_head + split_id * plse_stride_split + ) + lse = tl.load(partial_lse_ptr + plse_offset) + global_max_lse = tl.maximum(global_max_lse, lse) + + # Guard: if all splits had -inf LSE (empty sequence), output zeros + o_offset = batch_id * o_stride_batch + head_id * o_stride_head + dhead_offsets + if global_max_lse == float("-inf"): + tl.store(o_ptr + o_offset, tl.zeros([HEAD_DIM], dtype=tl.float32)) + return + + # Weighted combination: weight_i = exp(lse_i - global_max) + acc = tl.zeros([HEAD_DIM], dtype=tl.float32) + total_weight = 0.0 + + for split_id in range(NUM_SPLITS): + plse_offset = ( + batch_id * plse_stride_batch + head_id * plse_stride_head + split_id * plse_stride_split + ) + lse = tl.load(partial_lse_ptr + plse_offset) + weight = tl.exp(lse - global_max_lse) + + po_base = batch_id * po_stride_batch + head_id * po_stride_head + split_id * po_stride_split + partial_o = tl.load(partial_o_ptr + po_base + dhead_offsets) + + acc += weight * partial_o + total_weight += weight + + # Normalize and store + total_weight = tl.where(total_weight == 0.0, 1.0, total_weight) + o = acc / total_weight + tl.store(o_ptr + o_offset, o) + + +def triton_decode( + q: torch.Tensor, + kv_cache: torch.Tensor, + kv_indices: torch.Tensor, + kv_indptr: torch.Tensor, + kv_last_page_len: torch.Tensor, + sm_scale: float, + sliding_window: Optional[int] = None, + out: Optional[torch.Tensor] = None, +) -> torch.Tensor: + """Paged decode with GQA batching + FlashDecoding + page-aligned iteration. + + Args: + q: Query tensor [batch_size, n_heads, head_dim] + kv_cache: Combined cache [num_blocks, 2, n_kv_heads, page_size, head_dim] + kv_indices: Physical page indices (flattened) + kv_indptr: Cumulative page counts [batch_size + 1] + kv_last_page_len: Valid tokens in last page [batch_size] + sm_scale: Softmax scale factor + sliding_window: If set, only attend to the last sliding_window tokens + out: Optional output tensor [batch_size, n_heads, head_dim] + + Returns: + Output tensor [batch_size, n_heads, head_dim] + """ + batch_size, n_heads, head_dim = q.shape + _, _, n_kv_heads, page_size, _ = kv_cache.shape + head_ratio = n_heads // n_kv_heads + head_ratio_padded = max(1, 2 ** math.ceil(math.log2(head_ratio))) if head_ratio > 1 else 1 + page_block = _get_page_block(page_size) + + max_pages = kv_indices.shape[0] + max_seq_len = max_pages * page_size + # Normalize sliding_window: None/non-positive -> 0 (full attention) + sw = sliding_window if isinstance(sliding_window, int) and sliding_window > 0 else 0 + + output = out if out is not None else torch.empty_like(q) + + if batch_size == 0: + return output + + # Use effective sequence length (capped by sliding window) for split-K heuristic + effective_seq_len = min(max_seq_len, sw) if sw > 0 else max_seq_len + num_splits = _get_num_splits(effective_seq_len, batch_size, n_kv_heads, page_size) + + # Allocate intermediate buffers for split-K + partial_o = torch.empty( + batch_size, + n_heads, + num_splits, + head_dim, + dtype=torch.float32, + device=q.device, + ) + partial_lse = torch.empty( + batch_size, + n_heads, + num_splits, + dtype=torch.float32, + device=q.device, + ) + + # Stage 1: GQA-batched parallel KV processing + _flash_decode_stage1_kernel[(batch_size, n_kv_heads, num_splits)]( + q, + kv_cache, + kv_indices, + kv_indptr, + kv_last_page_len, + partial_o, + partial_lse, + # Q strides + q.stride(0), + q.stride(1), + # Partial output strides + partial_o.stride(0), + partial_o.stride(1), + partial_o.stride(2), + # Partial LSE strides + partial_lse.stride(0), + partial_lse.stride(1), + partial_lse.stride(2), + # Cache strides + kv_cache.stride(0), + kv_cache.stride(1), + kv_cache.stride(2), + kv_cache.stride(3), + # Constants + SM_SCALE=sm_scale, + N_HEADS=n_heads, + N_KV_HEADS=n_kv_heads, + HEAD_DIM=head_dim, + PAGE_SIZE=page_size, + PAGE_BLOCK=page_block, + HEAD_RATIO=head_ratio, + HEAD_RATIO_PADDED=head_ratio_padded, + NUM_SPLITS=num_splits, + SLIDING_WINDOW=sw, + ) + + # Stage 2: Combine partial results + _flash_decode_stage2_kernel[(batch_size, n_heads)]( + partial_o, + partial_lse, + output, + # Partial output strides + partial_o.stride(0), + partial_o.stride(1), + partial_o.stride(2), + # Partial LSE strides + partial_lse.stride(0), + partial_lse.stride(1), + partial_lse.stride(2), + # Output strides + output.stride(0), + output.stride(1), + # Constants + HEAD_DIM=head_dim, + NUM_SPLITS=num_splits, + ) + + return output diff --git a/tests/integration/test_lists/test-db/l0_h100.yml b/tests/integration/test_lists/test-db/l0_h100.yml index 7b26114e6126..3845d6d2c181 100644 --- a/tests/integration/test_lists/test-db/l0_h100.yml +++ b/tests/integration/test_lists/test-db/l0_h100.yml @@ -61,6 +61,15 @@ l0_h100: - unittest/_torch/modeling -k "modeling_gemma3" - unittest/_torch/modeling -k "modeling_gpt_oss" - unittest/_torch/modeling -k "modeling_whisper" # CPU-only log-mel parity + # Gemma4's head_dim=512 full-attention layers have no trtllm-gen cubins + # outside datacenter Blackwell and no FlashInfer paged kernel, so they run on + # the Triton prefill/decode path here. Individual node ids rather than the + # whole file: the rest of test_modeling_gemma4.py assumes trtllm-gen. + - unittest/_torch/modeling/test_modeling_gemma4.py::TestGemma4HFComparison::test_e2b_real_dims_triton_path + - unittest/_torch/modeling/test_modeling_gemma4.py::TestGemma4HFComparison::test_31b_real_dims_triton_path + - unittest/_torch/modeling/test_modeling_gemma4.py::TestGemma4HFComparison::test_26b_real_dims_triton_path + - unittest/_torch/modeling/test_modeling_gemma4.py::TestGemma4HFComparison::test_triton_path_is_actually_taken + - unittest/_torch/modeling/test_modeling_gemma4.py::TestGemma4ModelDefaults::test_only_full_attention_layers_exceed_flashinfer_head_dim_limit - unittest/_torch/modeling/test_modeling_nemotron_h.py::test_nemotron_h_sanity - unittest/_torch/modeling/test_multimodal_encoder_graph.py # Qwen3.5-MoE-VL is hybrid (Mamba SSM + attention); FlashInfer's diff --git a/tests/unittest/_torch/attention/backend_capability.py b/tests/unittest/_torch/attention/backend_capability.py index 19d9d7230dee..738b3e5afc1e 100644 --- a/tests/unittest/_torch/attention/backend_capability.py +++ b/tests/unittest/_torch/attention/backend_capability.py @@ -67,7 +67,14 @@ ), } -_FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS = (96, 512) +_FLASHINFER_PAGED_UNSUPPORTED_HEAD_DIMS = (96,) + +# head_dim>256 has no FlashInfer paged kernel. Below sm100 the backend routes +# these layers through Triton instead (the ``use_triton_attention`` path in +# FlashInferAttention.forward_impl), so they are testable there. On sm100+ that +# path is off -- the model-level backend uses trtllm-gen cubins, which this +# harness does not select (it builds FlashInfer with the default "fa2"). +_FLASHINFER_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS = (512,) _TRTLLM_PAGED_UNSUPPORTED_HEAD_DIMS = (512,) _TRTLLM_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS = (96,) @@ -133,6 +140,18 @@ def unsupported_reason(backend: str, case) -> Optional[str]: ): return f"FLASHINFER paged kernels do not support head_dim {case.head_dim}" + if ( + backend == "FLASHINFER" + and sm >= 100 + and getattr(case, "cache", "paged") != "none" + and not getattr(case, "is_mla", False) + and case.head_dim in _FLASHINFER_BLACKWELL_PAGED_UNSUPPORTED_HEAD_DIMS + ): + return ( + f"FLASHINFER head_dim {case.head_dim} needs trtllm-gen on sm100+; " + "the Triton fallback is only enabled below sm100" + ) + # TRTLLM's standard paged FMHA/MMHA kernels in this build do not cover the # Gemma4 head_dim 512 path. if ( diff --git a/tests/unittest/_torch/attention/test_triton_decode.py b/tests/unittest/_torch/attention/test_triton_decode.py new file mode 100644 index 000000000000..2f0c34ee234c --- /dev/null +++ b/tests/unittest/_torch/attention/test_triton_decode.py @@ -0,0 +1,388 @@ +# SPDX-FileCopyrightText: Copyright (c) 2022-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +"""Tests for the Triton FlashDecoding kernel with paged KV cache. + +Covers the shapes the PyTorch backend routes to Triton because FlashInfer's +paged kernels cannot serve them, in particular Gemma4's head_dim=512 +full-attention layers on architectures without trtllm-gen cubins. +""" + +import math + +import pytest +import torch + + +def _import_triton_decode(): + """Import triton_decode directly to avoid TRT-LLM C++ bindings.""" + import importlib.util + import os + + path = os.path.join( + os.path.dirname(__file__), + "..", + "..", + "..", + "..", + "tensorrt_llm", + "_torch", + "attention_backend", + "triton_decode.py", + ) + spec = importlib.util.spec_from_file_location("triton_decode", os.path.abspath(path)) + mod = importlib.util.module_from_spec(spec) + spec.loader.exec_module(mod) + return mod + + +triton_decode_mod = _import_triton_decode() +triton_decode = triton_decode_mod.triton_decode + +pytestmark = pytest.mark.skipif( + not torch.cuda.is_available(), reason="Triton decode kernel requires a GPU" +) + + +def _decode_reference(q, k_list, v_list, sm_scale, sliding_window=None): + """Reference single-token decode attention. + + Args: + q: [batch, n_heads, head_dim] + k_list: per-sequence [seq_len, n_kv_heads, head_dim] + v_list: per-sequence [seq_len, n_kv_heads, head_dim] + sm_scale: softmax scale + sliding_window: attend only to the last N tokens (inclusive of current) + + Returns: + [batch, n_heads, head_dim] + """ + batch, n_heads, _ = q.shape + outputs = [] + for b in range(batch): + k = k_list[b] + v = v_list[b] + seq_len, n_kv_heads, _ = k.shape + if sliding_window is not None and sliding_window > 0: + start = max(0, seq_len - sliding_window) + k = k[start:] + v = v[start:] + gqa = n_heads // n_kv_heads + # Broadcast KV heads across their query group. + k_rep = k.repeat_interleave(gqa, dim=1) # [win, n_heads, head_dim] + v_rep = v.repeat_interleave(gqa, dim=1) + # [n_heads, win] + scores = torch.einsum("hd,whd->hw", q[b].float(), k_rep.float()) + scores = scores * sm_scale + probs = torch.softmax(scores, dim=-1) + out = torch.einsum("hw,whd->hd", probs, v_rep.float()) + outputs.append(out) + return torch.stack(outputs, dim=0) + + +def _build_paged_cache(k_list, v_list, page_size, device, dtype=torch.bfloat16): + """Pack per-sequence KV into a combined HND paged cache. + + Returns: + kv_cache: [num_pages, 2, n_kv_heads, page_size, head_dim] + kv_indices: [total_pages] physical page ids + kv_indptr: [batch + 1] cumulative page counts + kv_last_page_len: [batch] valid tokens in each sequence's last page + """ + n_kv_heads, head_dim = k_list[0].shape[1], k_list[0].shape[2] + pages_k, pages_v = [], [] + page_counts = [0] + last_page_lens = [] + + for k, v in zip(k_list, v_list): + seq_len = k.shape[0] + n_pages = (seq_len + page_size - 1) // page_size + page_counts.append(page_counts[-1] + n_pages) + last = seq_len - (n_pages - 1) * page_size + last_page_lens.append(last) + + for p in range(n_pages): + start = p * page_size + end = min(start + page_size, seq_len) + n_tok = end - start + page_k = torch.zeros(n_kv_heads, page_size, head_dim, dtype=dtype, device=device) + page_v = torch.zeros(n_kv_heads, page_size, head_dim, dtype=dtype, device=device) + page_k[:, :n_tok, :] = k[start:end].transpose(0, 1).to(dtype) + page_v[:, :n_tok, :] = v[start:end].transpose(0, 1).to(dtype) + pages_k.append(page_k) + pages_v.append(page_v) + + k_pages = torch.stack(pages_k, dim=0) + v_pages = torch.stack(pages_v, dim=0) + kv_cache = torch.stack([k_pages, v_pages], dim=1) + + kv_indices = torch.arange(k_pages.shape[0], dtype=torch.int32, device=device) + kv_indptr = torch.tensor(page_counts, dtype=torch.int32, device=device) + kv_last_page_len = torch.tensor(last_page_lens, dtype=torch.int32, device=device) + return kv_cache, kv_indices, kv_indptr, kv_last_page_len + + +def _run( + seq_lens, + n_heads, + n_kv_heads, + head_dim, + page_size, + device, + sliding_window=None, + kv_dtype=torch.bfloat16, + seed=0, +): + """Build inputs, run the kernel, and return (actual, expected).""" + torch.manual_seed(seed) + batch = len(seq_lens) + + q = torch.randn(batch, n_heads, head_dim, dtype=torch.bfloat16, device=device) + k_list = [ + torch.randn(s, n_kv_heads, head_dim, dtype=torch.bfloat16, device=device) for s in seq_lens + ] + v_list = [ + torch.randn(s, n_kv_heads, head_dim, dtype=torch.bfloat16, device=device) for s in seq_lens + ] + + kv_cache, kv_indices, kv_indptr, kv_last_page_len = _build_paged_cache( + k_list, v_list, page_size, device, dtype=kv_dtype + ) + + # The kernel reads the cache and casts to the query dtype, so an FP8 cache + # is lossy. Compare against the values actually stored, not the originals. + if kv_dtype != torch.bfloat16: + k_list = [k.to(kv_dtype).to(torch.bfloat16) for k in k_list] + v_list = [v.to(kv_dtype).to(torch.bfloat16) for v in v_list] + + sm_scale = 1.0 / math.sqrt(head_dim) + actual = triton_decode( + q=q, + kv_cache=kv_cache, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=kv_last_page_len, + sm_scale=sm_scale, + sliding_window=sliding_window, + ) + expected = _decode_reference(q, k_list, v_list, sm_scale, sliding_window=sliding_window) + return actual.float(), expected + + +@pytest.fixture +def device(): + return torch.device("cuda:0") + + +class TestTritonDecodeHeadDims: + """head_dim coverage, including the Gemma4 512 path.""" + + @pytest.mark.parametrize("head_dim", [64, 128, 256, 512]) + def test_head_dims(self, head_dim, device): + actual, expected = _run( + [37], n_heads=8, n_kv_heads=4, head_dim=head_dim, page_size=16, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def test_gemma4_e2b_mqa_hd512(self, device): + """Gemma4-E2B full-attention layers: MQA, head_dim 512.""" + actual, expected = _run( + [64], n_heads=8, n_kv_heads=1, head_dim=512, page_size=16, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def test_gemma4_31b_gqa_hd512(self, device): + """Gemma4-31B full-attention layers: GQA 32/4, head_dim 512.""" + actual, expected = _run( + [100], n_heads=32, n_kv_heads=4, head_dim=512, page_size=32, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +class TestTritonDecodeGQARatios: + """GQA/MQA head-ratio coverage, including non-power-of-2 ratios.""" + + @pytest.mark.parametrize( + "n_heads,n_kv_heads", + [ + (8, 8), + (8, 4), + (8, 1), + (16, 2), + (12, 4), + (24, 4), + ], + ) + def test_head_ratios(self, n_heads, n_kv_heads, device): + actual, expected = _run( + [48], n_heads=n_heads, n_kv_heads=n_kv_heads, head_dim=128, page_size=16, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +class TestTritonDecodePaging: + """Page-boundary and partial-last-page behaviour.""" + + @pytest.mark.parametrize("page_size", [1, 8, 16, 32, 64]) + def test_page_sizes(self, page_size, device): + actual, expected = _run( + [53], n_heads=8, n_kv_heads=2, head_dim=128, page_size=page_size, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def test_exact_page_multiple(self, device): + """seq_len an exact multiple of page_size: last page is full.""" + actual, expected = _run( + [64], n_heads=8, n_kv_heads=2, head_dim=128, page_size=16, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def test_single_token_sequence(self, device): + actual, expected = _run( + [1], n_heads=8, n_kv_heads=2, head_dim=128, page_size=16, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def test_variable_length_batch(self, device): + """Mixed sequence lengths exercise per-sequence indptr/last_page_len.""" + actual, expected = _run( + [1, 15, 16, 17, 128, 300], + n_heads=8, + n_kv_heads=2, + head_dim=128, + page_size=16, + device=device, + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def test_long_sequence_multi_split(self, device): + """Long single sequence with small batch forces split-K > 1.""" + actual, expected = _run( + [4096], n_heads=8, n_kv_heads=1, head_dim=128, page_size=32, device=device + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +class TestTritonDecodeSlidingWindow: + """Sliding-window masking, as used by Gemma4's VSWA layers.""" + + @pytest.mark.parametrize("sliding_window", [1, 8, 64, 1024]) + def test_sliding_window(self, sliding_window, device): + actual, expected = _run( + [300], + n_heads=8, + n_kv_heads=2, + head_dim=128, + page_size=16, + device=device, + sliding_window=sliding_window, + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + def test_window_larger_than_sequence(self, device): + """A window wider than the sequence must equal full attention.""" + windowed, _ = _run( + [32], + n_heads=8, + n_kv_heads=2, + head_dim=128, + page_size=16, + device=device, + sliding_window=4096, + seed=7, + ) + full, _ = _run( + [32], + n_heads=8, + n_kv_heads=2, + head_dim=128, + page_size=16, + device=device, + sliding_window=None, + seed=7, + ) + torch.testing.assert_close(windowed, full, atol=2e-2, rtol=2e-2) + + def test_sliding_window_hd512_batch(self, device): + actual, expected = _run( + [120, 500], + n_heads=8, + n_kv_heads=1, + head_dim=512, + page_size=32, + device=device, + sliding_window=1024, + ) + torch.testing.assert_close(actual, expected, atol=2e-2, rtol=2e-2) + + +class TestTritonDecodeFp8Cache: + """FP8 KV cache is dequantized in-kernel and needs no conversion pass.""" + + @pytest.mark.parametrize("head_dim", [128, 512]) + def test_fp8_kv_cache(self, head_dim, device): + actual, expected = _run( + [64], + n_heads=8, + n_kv_heads=2, + head_dim=head_dim, + page_size=16, + device=device, + kv_dtype=torch.float8_e4m3fn, + ) + # Looser bound: e4m3 round-trip dominates the error. + torch.testing.assert_close(actual, expected, atol=2e-1, rtol=2e-1) + + def test_fp8_sliding_window(self, device): + actual, expected = _run( + [256], + n_heads=8, + n_kv_heads=2, + head_dim=128, + page_size=16, + device=device, + sliding_window=64, + kv_dtype=torch.float8_e4m3fn, + ) + torch.testing.assert_close(actual, expected, atol=2e-1, rtol=2e-1) + + +class TestTritonDecodeOutputTensor: + """The caller-supplied ``out`` tensor is written in place.""" + + def test_out_written_in_place(self, device): + torch.manual_seed(3) + n_heads, n_kv_heads, head_dim, page_size = 8, 2, 128, 16 + q = torch.randn(1, n_heads, head_dim, dtype=torch.bfloat16, device=device) + k_list = [torch.randn(40, n_kv_heads, head_dim, dtype=torch.bfloat16, device=device)] + v_list = [torch.randn(40, n_kv_heads, head_dim, dtype=torch.bfloat16, device=device)] + kv_cache, kv_indices, kv_indptr, last_page = _build_paged_cache( + k_list, v_list, page_size, device + ) + + out = torch.zeros_like(q) + sm_scale = 1.0 / math.sqrt(head_dim) + returned = triton_decode( + q=q, + kv_cache=kv_cache, + kv_indices=kv_indices, + kv_indptr=kv_indptr, + kv_last_page_len=last_page, + sm_scale=sm_scale, + out=out, + ) + + assert returned.data_ptr() == out.data_ptr() + expected = _decode_reference(q, k_list, v_list, sm_scale) + torch.testing.assert_close(out.float(), expected, atol=2e-2, rtol=2e-2) diff --git a/tests/unittest/_torch/modeling/test_modeling_gemma4.py b/tests/unittest/_torch/modeling/test_modeling_gemma4.py index 5eff2c644555..b3f65a5ec6ba 100644 --- a/tests/unittest/_torch/modeling/test_modeling_gemma4.py +++ b/tests/unittest/_torch/modeling/test_modeling_gemma4.py @@ -1220,6 +1220,80 @@ def test_kev_config(self): """K=V attention: full layers share key→value, v_norm applied.""" self._run_full_model_comparison(deepcopy(GEMMA4_KEV_CONFIG)) + # -- Production head dims (256 sliding / 512 full) without trtllm-gen ---- + # + # trtllm-gen ships head_dim 512 cubins for datacenter Blackwell only, and + # FlashInfer's fa2/fa3 paged kernels stop at 256. Patching ``is_sm_100f`` + # off forces the dispatch every other architecture takes (Hopper included), + # where the full-attention layers run on the Triton prefill/decode kernels + # while the sliding layers stay on FlashInfer. Comparing against HF then + # covers the mixed-backend model end to end, on any GPU. + + def _run_triton_path_comparison(self, config_dict): + """Run the HF comparison with the non-Blackwell dispatch forced on. + + Both ``is_sm_100f`` call sites must be patched together. The one in + ``modeling_gemma4`` picks ``flashinfer_backend`` ("fa2" off Blackwell, + "trtllm-gen" on it); the one in the FlashInfer backend decides whether + head_dim>256 layers divert to Triton. Patching only the backend would, + on Blackwell, leave the sliding layers on trtllm-gen while the full + layers ran on Triton -- a mix that never occurs in production. Patching + both makes this a faithful non-Blackwell simulation on any GPU. + """ + with ( + unittest.mock.patch( + "tensorrt_llm._torch.attention_backend.flashinfer.is_sm_100f", + return_value=False, + ), + unittest.mock.patch( + "tensorrt_llm._torch.models.modeling_gemma4.is_sm_100f", + return_value=False, + ), + ): + self._run_full_model_comparison(config_dict) + + @torch.no_grad() + def test_e2b_real_dims_triton_path(self): + """E2B geometry: MQA sliding + head_dim 512 full layers on Triton.""" + self._run_triton_path_comparison(deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG)) + + @torch.no_grad() + def test_31b_real_dims_triton_path(self): + """31B geometry: mixed GQA ratios and K=V full layers on Triton.""" + self._run_triton_path_comparison(deepcopy(GEMMA4_31B_REAL_DIMS_CONFIG)) + + @torch.no_grad() + def test_26b_real_dims_triton_path(self): + """26B-A4B geometry: GQA 16/8 sliding and K=V full layers on Triton.""" + self._run_triton_path_comparison(deepcopy(GEMMA4_26B_REAL_DIMS_CONFIG)) + + @torch.no_grad() + def test_triton_path_is_actually_taken(self): + """Both Triton kernels must really run for the head_dim 512 layers. + + Without this, a dispatch regression that quietly sent head_dim 512 back + to FlashInfer would only surface as a kernel error on some + architectures, and the comparisons above would look like passing + coverage of a path they never executed. + """ + import tensorrt_llm._torch.attention_backend.triton_decode as decode_mod + import tensorrt_llm._torch.attention_backend.triton_prefill as prefill_mod + + with ( + unittest.mock.patch.object( + decode_mod, "triton_decode", wraps=decode_mod.triton_decode + ) as decode_spy, + unittest.mock.patch.object( + prefill_mod, + "triton_prefill_with_custom_mask", + wraps=prefill_mod.triton_prefill_with_custom_mask, + ) as prefill_spy, + ): + self._run_triton_path_comparison(deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG)) + + self.assertTrue(prefill_spy.called, "head_dim 512 prefill never reached Triton") + self.assertTrue(decode_spy.called, "head_dim 512 decode never reached Triton") + @torch.no_grad() def test_kev_vnorm_order(self): """K=V: v_norm must apply to raw k_proj(x), NOT after k_norm. @@ -2454,6 +2528,43 @@ def test_non_sm100f_layers_use_fa2(self, _mock_is_sm_100f): ) self.assertEqual(attn.attn.flashinfer_backend, "fa2") + def test_only_full_attention_layers_exceed_flashinfer_head_dim_limit(self): + """Production head dims decide which layers FlashInfer can serve. + + ``flashinfer_backend == "fa2"`` off Blackwell does not mean fa2 runs + every layer: ``FlashInferAttention.forward_impl`` sends anything above + ``_FLASHINFER_MAX_PAGED_HEAD_DIM`` to Triton instead, because fa2/fa3 + have no paged kernel for it. Pin the geometry that split depends on -- + Gemma4's sliding layers sit at the limit and its full-attention layers + sit above it, so a change to either would silently reroute layers. + """ + from tensorrt_llm._torch.attention_backend.flashinfer import _FLASHINFER_MAX_PAGED_HEAD_DIM + + config = Gemma4TextConfig(**deepcopy(GEMMA4_E2B_REAL_DIMS_CONFIG)) + seen = set() + for layer_type in config.layer_types: + is_sliding = layer_type == "sliding_attention" + head_dim = config.head_dim if is_sliding else config.global_head_dim + seen.add(layer_type) + if is_sliding: + self.assertLessEqual( + head_dim, + _FLASHINFER_MAX_PAGED_HEAD_DIM, + "sliding layers are expected to stay on FlashInfer", + ) + else: + self.assertGreater( + head_dim, + _FLASHINFER_MAX_PAGED_HEAD_DIM, + "full-attention layers are expected to need the Triton path", + ) + + self.assertEqual( + seen, + {"sliding_attention", "full_attention"}, + "config must exercise both sides of the head_dim split", + ) + class TestGemma4CUDAGraph(unittest.TestCase): """Tests for Gemma4 attention with CUDA graph capture/replay."""