[Fix][trtllm-mha] Canonicalize the strides when num_head = 1 (#17732)

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.

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)