Fix trtllm mla chunked-prefill zero-length bug (#22291) (#22688)

python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh

@@ -0,0 +1,124 @@
+#pragma once
+
+// Fixup kernel for TRT-LLM ragged attention zero-KV rows.
+// For sequences with kv_len == 0, forces out=0 and lse=-inf.
+// 2D grid: (blocks_per_seq, batch_size). Y-dim early-exits for non-zero KV.
+// Uses vectorised float4 stores for bandwidth efficiency.
+
+#include <sgl_kernel/tensor.h>
+
+#include <sgl_kernel/utils.cuh>
+
+#include <cstdint>
+
+namespace {
+
+constexpr int kFixupBlockSize = 256;
+
+// -- vectorised zero-fill helpers ------------------------------------------
+
+// Zero-fill `n` elements of type T starting at `ptr`, using float4 stores.
+// `ptr` must be 16-byte aligned (guaranteed by PyTorch allocator).
+template <typename T>
+__device__ __forceinline__ void vec_zero_fill(T* ptr, int n) {
+  constexpr int kVec = 16 / sizeof(T);  // elements per float4
+  const int n_vec = n / kVec;           // full vectors
+  float4* dst4 = reinterpret_cast<float4*>(ptr);
+  const float4 z4 = make_float4(0.f, 0.f, 0.f, 0.f);
+  for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
+    dst4[i] = z4;
+  }
+  // tail elements
+  const int tail_start = n_vec * kVec;
+  for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
+    ptr[i] = static_cast<T>(0);
+  }
+}
+
+// Fill `n` float elements with -inf using float4 stores.
+__device__ __forceinline__ void vec_neginf_fill(float* ptr, int n) {
+  constexpr int kVec = 4;  // float4 = 4 floats
+  const int n_vec = n / kVec;
+  float4* dst4 = reinterpret_cast<float4*>(ptr);
+  const float ninf = -INFINITY;
+  const float4 inf4 = make_float4(ninf, ninf, ninf, ninf);
+  for (int i = threadIdx.x; i < n_vec; i += blockDim.x) {
+    dst4[i] = inf4;
+  }
+  const int tail_start = n_vec * kVec;
+  for (int i = tail_start + threadIdx.x; i < n; i += blockDim.x) {
+    ptr[i] = ninf;
+  }
+}
+
+// -- main kernel -----------------------------------------------------------
+
+template <typename OutT>
+__global__ void fixup_zero_kv_rows_kernel(
+    OutT* __restrict__ out,
+    float* __restrict__ lse,
+    const int32_t* __restrict__ kv_lens,
+    const int32_t* __restrict__ cum_seq_lens,
+    const int out_stride,
+    const int lse_stride) {
+  const int seq_idx = blockIdx.y;
+  if (kv_lens[seq_idx] > 0) return;
+
+  const int tok_start = cum_seq_lens[seq_idx];
+  const int tok_end = cum_seq_lens[seq_idx + 1];
+  const int num_tokens = tok_end - tok_start;
+  if (num_tokens <= 0) return;
+
+  // blockIdx.x selects a token within this sequence.
+  const int tok = tok_start + blockIdx.x;
+  if (tok >= tok_end) return;
+
+  // Each block handles one token: zero out[tok] and set lse[tok] = -inf.
+  vec_zero_fill(out + tok * out_stride, out_stride);
+  vec_neginf_fill(lse + tok * lse_stride, lse_stride);
+}
+
+// -- host launcher ---------------------------------------------------------
+
+template <typename OutT>
+void fixup_zero_kv_rows(
+    tvm::ffi::TensorView out,
+    tvm::ffi::TensorView lse,
+    tvm::ffi::TensorView kv_lens,
+    tvm::ffi::TensorView cum_seq_lens,
+    int64_t max_seq_len) {
+  using namespace host;
+
+  auto batch_size = SymbolicSize{"batch_size"};
+  auto total_tokens = SymbolicSize{"total_tokens"};
+  auto num_heads = SymbolicSize{"num_heads"};
+  auto v_head_dim = SymbolicSize{"v_head_dim"};
+  auto batch_size_plus_1 = SymbolicSize{"batch_size_plus_1"};
+  auto device = SymbolicDevice{};
+  device.set_options<kDLCUDA>();
+
+  TensorMatcher({total_tokens, num_heads, v_head_dim}).with_dtype<OutT>().with_device(device).verify(out);
+  TensorMatcher({total_tokens, num_heads}).with_dtype<float>().with_device(device).verify(lse);
+  TensorMatcher({batch_size}).with_dtype<int32_t>().with_device(device).verify(kv_lens);
+  TensorMatcher({batch_size_plus_1}).with_dtype<int32_t>().with_device(device).verify(cum_seq_lens);
+
+  const int bs = static_cast<int>(batch_size.unwrap());
+  const int nh = static_cast<int>(num_heads.unwrap());
+  const int vd = static_cast<int>(v_head_dim.unwrap());
+
+  // Grid: one block per (token, sequence). X = max tokens in any seq.
+  const int blocks_x = static_cast<int>(max_seq_len);
+  dim3 grid(blocks_x, bs);
+  dim3 block(kFixupBlockSize);
+
+  LaunchKernel(grid, block, device.unwrap())(
+      fixup_zero_kv_rows_kernel<OutT>,
+      static_cast<OutT*>(out.data_ptr()),
+      static_cast<float*>(lse.data_ptr()),
+      static_cast<const int32_t*>(kv_lens.data_ptr()),
+      static_cast<const int32_t*>(cum_seq_lens.data_ptr()),
+      nh * vd,
+      nh);
+}
+
+}  // namespace

python/sglang/jit_kernel/fixup_zero_kv.py

@@ -0,0 +1,44 @@
+from __future__ import annotations
+
+from typing import TYPE_CHECKING
+
+import torch
+
+from sglang.jit_kernel.utils import cache_once, load_jit, make_cpp_args
+
+if TYPE_CHECKING:
+    from tvm_ffi.module import Module
+
+
+@cache_once
+def _jit_fixup_module(dtype: torch.dtype) -> Module:
+    args = make_cpp_args(dtype)
+    return load_jit(
+        "fixup_zero_kv",
+        *args,
+        cuda_files=["attention/fixup_zero_kv.cuh"],
+        cuda_wrappers=[("fixup_zero_kv_rows", f"fixup_zero_kv_rows<{args}>")],
+    )
+
+
+def fixup_zero_kv_rows(
+    out: torch.Tensor,
+    lse: torch.Tensor,
+    kv_lens: torch.Tensor,
+    cum_seq_lens: torch.Tensor,
+    max_seq_len: int,
+) -> None:
+    """Fix output and LSE for zero-KV rows after TRT-LLM ragged attention.
+
+    For sequences with kv_lens[i] == 0, sets out[tokens_i] = 0 and
+    lse[tokens_i] = -inf.  Single CUDA kernel launch, no GPU-CPU sync.
+
+    Args:
+        out:          [total_tokens, num_heads, v_head_dim]  bf16/fp16
+        lse:          [total_tokens, num_heads]               float32
+        kv_lens:      [batch_size]                            int32
+        cum_seq_lens: [batch_size + 1]                        int32
+        max_seq_len:  max Q tokens in any single sequence     int
+    """
+    module = _jit_fixup_module(out.dtype)
+    module.fixup_zero_kv_rows(out, lse, kv_lens, cum_seq_lens, max_seq_len)

python/sglang/srt/layers/attention/trtllm_mla_backend.py

@@ -13,6 +13,7 @@ import torch
 import triton
 import triton.language as tl
 
+from sglang.jit_kernel.fixup_zero_kv import fixup_zero_kv_rows
 from sglang.srt.compilation.piecewise_context_manager import is_in_piecewise_cuda_graph
 from sglang.srt.environ import envs
 from sglang.srt.layers.attention.flashinfer_mla_backend import (
@@ -1121,7 +1122,7 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 dtype=self.q_data_type,
                 device=q.device,
             )
-            return flashinfer.prefill.trtllm_ragged_attention_deepseek(
+            result = flashinfer.prefill.trtllm_ragged_attention_deepseek(
                 **common_trtllm_args,
                 seq_lens=forward_batch.prefix_chunk_seq_lens[chunk_idx],
                 max_kv_len=forward_batch.prefix_chunk_max_seq_lens[chunk_idx],
@@ -1131,6 +1132,25 @@ class TRTLLMMLABackend(FlashInferMLAAttnBackend):
                 return_lse=True,
                 out=out,
             )
+
+            # The TRT-LLM ragged attention cubin kernel does not correctly
+            # handle rows with kv_len == 0: it leaves stale data in the
+            # workspace softmaxStats buffer and may produce non-zero output
+            # for those rows.  Fix up by forcing out=0 and lse=-inf for
+            # zero-KV rows so that downstream merge_state ignores them.
+            # Skip entirely when this chunk has no zero-KV rows (pure CPU
+            # check, precomputed in prepare_chunked_prefix_cache_info).
+            if forward_batch.prefix_chunk_has_zero_kv[chunk_idx]:
+                out_tensor, lse_tensor = result
+                fixup_zero_kv_rows(
+                    out_tensor,
+                    lse_tensor,
+                    forward_batch.prefix_chunk_seq_lens[chunk_idx],
+                    self.forward_prefill_metadata.cum_seq_lens,
+                    self.forward_prefill_metadata.max_seq_len,
+                )
+
+            return result
         else:
             out = torch.zeros(
                 q.shape[0],

python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py

@@ -28,6 +28,9 @@ class ForwardBatchDeepSeekMHAMixin:
     prefix_chunk_cu_seq_lens: Optional[torch.Tensor] = None
     # Max lengths of prefix cache for each chunk, (num_prefix_chunks,)
     prefix_chunk_max_seq_lens: Optional[List[int]] = None
+    # Per-chunk flag: True if any sequence has kv_len==0 in that chunk.
+    # Precomputed on CPU to avoid GPU-CPU sync in the hot path.
+    prefix_chunk_has_zero_kv: Optional[List[bool]] = None
     # Number of tokens in each prefix cache chunk, (num_prefix_chunks,)
     prefix_chunk_num_tokens: Optional[List[int]] = None
     # KV Indices for each chunk
@@ -163,6 +166,13 @@ class ForwardBatchDeepSeekMHAMixin:
         self.prefix_chunk_num_tokens = prefix_chunk_seq_lens_cpu.sum(dim=1).tolist()
         assert max(self.prefix_chunk_num_tokens) <= self.get_max_chunk_capacity()
 
+        # Per-chunk flag: does any sequence have kv_len == 0?
+        # Pure CPU check (prefix_chunk_seq_lens_cpu is on CPU), no GPU sync.
+        self.prefix_chunk_has_zero_kv = [
+            bool((prefix_chunk_seq_lens_cpu[i] == 0).any())
+            for i in range(self.num_prefix_chunks)
+        ]
+
         # Precompute the kv indices for each chunk
         self.prepare_chunked_kv_indices(device)