From 673dc09d9ba30331fa60f23c88c0483865d3afce Mon Sep 17 00:00:00 2001 From: Jerry Ji Date: Wed, 28 Jan 2026 18:11:18 -0800 Subject: [PATCH] [Fix][trtllm-mha] Canonicalize the strides when num_head = 1 (#17732) --- .../layers/attention/trtllm_mha_backend.py | 13 +++++++ python/sglang/srt/layers/attention/utils.py | 35 +++++++++++++++++++ 2 files changed, 48 insertions(+) diff --git a/python/sglang/srt/layers/attention/trtllm_mha_backend.py b/python/sglang/srt/layers/attention/trtllm_mha_backend.py index aa418b7af..4773a0d2a 100644 --- a/python/sglang/srt/layers/attention/trtllm_mha_backend.py +++ b/python/sglang/srt/layers/attention/trtllm_mha_backend.py @@ -19,6 +19,7 @@ from sglang.srt.layers.attention.flashinfer_backend import ( from sglang.srt.layers.attention.triton_ops.trtllm_fp8_kv_kernel import ( fused_fp8_set_kv_buffer, ) +from sglang.srt.layers.attention.utils import canonicalize_stride from sglang.srt.model_executor.forward_batch_info import ForwardBatch, ForwardMode from sglang.srt.utils import is_flashinfer_available @@ -608,6 +609,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): v_cache = v_cache.view( -1, self.page_size, layer.tp_v_head_num, layer.head_dim ).permute(0, 2, 1, 3) + + if layer.tp_k_head_num == 1: + k_cache = canonicalize_stride(k_cache) + if layer.tp_v_head_num == 1: + v_cache = canonicalize_stride(v_cache) + kv_cache = (k_cache, v_cache) # TODO: add support for quantization @@ -684,6 +691,12 @@ class TRTLLMHAAttnBackend(FlashInferAttnBackend): v_cache = v_cache.view( -1, self.page_size, layer.tp_v_head_num, layer.head_dim ).permute(0, 2, 1, 3) + + if layer.tp_k_head_num == 1: + k_cache = canonicalize_stride(k_cache) + if layer.tp_v_head_num == 1: + v_cache = canonicalize_stride(v_cache) + kv_cache = (k_cache, v_cache) # sink: additional value per head in the denominator of the softmax. diff --git a/python/sglang/srt/layers/attention/utils.py b/python/sglang/srt/layers/attention/utils.py index dc9973788..f01fa4094 100644 --- a/python/sglang/srt/layers/attention/utils.py +++ b/python/sglang/srt/layers/attention/utils.py @@ -277,3 +277,38 @@ def pad_sequence_with_mask( ) return B, output, attn_mask + + +# When num_kv_heads=1, we have tensors with degenerate strides, +# For example, as below, where we have stride[-3] == stride[-2]: +# - shape: [num_pages, 1, 64, 128] +# - stride: [8192, 128, 128, 1] +# This will cause TMA desc validation fail in flashinfer (trtllm-mha backend). +# +# See: https://github.com/flashinfer-ai/flashinfer/issues/2232 +def canonicalize_stride(tensor: torch.Tensor) -> torch.Tensor: + """ + Adjust degenerate strides for a tensor, make it canonical. + """ + sizes = tensor.size() + strides = tensor.stride() + ndim = tensor.dim() + + need_fix = any( + sizes[i] == 1 and strides[i] == strides[i + 1] for i in range(ndim - 1) + ) + + if not need_fix: + return tensor + + # canonicalize the stride + # Example: + # - shape: [num_pages, 1, 64, 128] + # - stride: [8192, 128, 128, 1] (wrong!) + # Gives new stride: [8192, 8192, 128 ,1] (correct!) + new_strides = [0] * ndim + new_strides[-1] = 1 + for i in range(ndim - 2, -1, -1): + new_strides[i] = new_strides[i + 1] * sizes[i + 1] + + return tensor.as_strided(sizes, new_strides)