From 5595f6e9888b8c8ee9e8cc627769131b94c27698 Mon Sep 17 00:00:00 2001 From: Yuhao Yang <47235274+yhyang201@users.noreply.github.com> Date: Tue, 21 Apr 2026 13:10:13 +0800 Subject: [PATCH] Fix trtllm mla chunked-prefill zero-length bug (#22291) (#22688) --- .../csrc/attention/fixup_zero_kv.cuh | 124 ++++++++++++++++++ python/sglang/jit_kernel/fixup_zero_kv.py | 44 +++++++ .../layers/attention/trtllm_mla_backend.py | 22 +++- .../forward_batch_deepseek_mha_mixin.py | 10 ++ 4 files changed, 199 insertions(+), 1 deletion(-) create mode 100644 python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh create mode 100644 python/sglang/jit_kernel/fixup_zero_kv.py diff --git a/python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh b/python/sglang/jit_kernel/csrc/attention/fixup_zero_kv.cuh new file mode 100644 index 000000000..5d1633e32 --- /dev/null +++ b/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 + +#include + +#include + +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 +__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(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(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(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 +__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 +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(); + + TensorMatcher({total_tokens, num_heads, v_head_dim}).with_dtype().with_device(device).verify(out); + TensorMatcher({total_tokens, num_heads}).with_dtype().with_device(device).verify(lse); + TensorMatcher({batch_size}).with_dtype().with_device(device).verify(kv_lens); + TensorMatcher({batch_size_plus_1}).with_dtype().with_device(device).verify(cum_seq_lens); + + const int bs = static_cast(batch_size.unwrap()); + const int nh = static_cast(num_heads.unwrap()); + const int vd = static_cast(v_head_dim.unwrap()); + + // Grid: one block per (token, sequence). X = max tokens in any seq. + const int blocks_x = static_cast(max_seq_len); + dim3 grid(blocks_x, bs); + dim3 block(kFixupBlockSize); + + LaunchKernel(grid, block, device.unwrap())( + fixup_zero_kv_rows_kernel, + static_cast(out.data_ptr()), + static_cast(lse.data_ptr()), + static_cast(kv_lens.data_ptr()), + static_cast(cum_seq_lens.data_ptr()), + nh * vd, + nh); +} + +} // namespace diff --git a/python/sglang/jit_kernel/fixup_zero_kv.py b/python/sglang/jit_kernel/fixup_zero_kv.py new file mode 100644 index 000000000..6175c0f37 --- /dev/null +++ b/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) diff --git a/python/sglang/srt/layers/attention/trtllm_mla_backend.py b/python/sglang/srt/layers/attention/trtllm_mla_backend.py index 772c3dac8..65c1cdb54 100755 --- a/python/sglang/srt/layers/attention/trtllm_mla_backend.py +++ b/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], diff --git a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py b/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py index 26a59a6cc..2840f9f23 100644 --- a/python/sglang/srt/model_executor/forward_batch_deepseek_mha_mixin.py +++ b/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)