From a74470e9046684d73a4492a5523ba5ee35729752 Mon Sep 17 00:00:00 2001 From: Xiaoyu Zhang <1182563586@qq.com> Date: Sat, 5 Sep 2026 17:08:20 +0800 Subject: [PATCH] fix(mamba): unify causal_conv1d col* dtype to x (MiniCPM-V-4.6 GDN prefill bf16/fp16 mismatch) (#38039) --- .../kernels/ops/mamba/causal_conv1d_triton.py | 30 ++++++---- .../layers/mamba/test_causal_conv1d.py | 56 +++++++++++++++++++ 2 files changed, 74 insertions(+), 12 deletions(-) diff --git a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py index 321a15d0d..9449a3a43 100644 --- a/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py +++ b/python/sglang/kernels/ops/mamba/causal_conv1d_triton.py @@ -123,35 +123,41 @@ def _causal_conv1d_fwd_kernel( # continuous batching if HAS_INITIAL_STATES: # the new HAS_INITIAL_STATES load_init_state = tl.load(has_initial_states_ptr + idx_seq).to(tl.int1) if load_init_state: - # load from conv_states + # load from conv_states. Cast to x's dtype so col* keep a single + # dtype across the whole kernel: when x is fp16 but the conv-state + # cache is bf16 (e.g. MiniCPM-V GDN prefill), the chunk_offset==0 + # branch would otherwise produce bf16 cols while the chunk_offset>0 + # else branch (and the sliding-window reassignment) produce fp16, + # which trips Triton's if/else phi type check on col0. + x_elem_ty = x_ptr.dtype.element_ty prior_tokens = conv_states_base + (state_len - 1) * stride_conv_state_tok mask_w = idx_feats < dim if KERNEL_WIDTH == 2: conv_states_ptrs = prior_tokens # [BLOCK_N] - col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + col0 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) if KERNEL_WIDTH == 3: conv_states_ptrs = prior_tokens # [BLOCK_N] - col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + col1 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] - col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + col0 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) if KERNEL_WIDTH == 4: conv_states_ptrs = prior_tokens # [BLOCK_N] - col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + col2 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] - col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + col1 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] - col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + col0 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) if KERNEL_WIDTH == 5: conv_states_ptrs = prior_tokens # [BLOCK_N] - col3 = tl.load(conv_states_ptrs, mask_w, 0.0) + col3 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) conv_states_ptrs = prior_tokens - 1 * stride_conv_state_tok # [BLOCK_N] - col2 = tl.load(conv_states_ptrs, mask_w, 0.0) + col2 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) conv_states_ptrs = prior_tokens - 2 * stride_conv_state_tok # [BLOCK_N] - col1 = tl.load(conv_states_ptrs, mask_w, 0.0) + col1 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) conv_states_ptrs = prior_tokens - 3 * stride_conv_state_tok # [BLOCK_N] - col0 = tl.load(conv_states_ptrs, mask_w, 0.0) + col0 = tl.load(conv_states_ptrs, mask_w, 0.0).to(x_elem_ty) else: - # prior-tokens are zeros + # prior-tokens are zeros (same x dtype as every other col* source) if KERNEL_WIDTH >= 2: # STRATEGY1 # first chunk and does not have prior-token, so just set to 0 col0 = tl.zeros((BLOCK_N,), dtype=x_ptr.dtype.element_ty) diff --git a/test/registered/layers/mamba/test_causal_conv1d.py b/test/registered/layers/mamba/test_causal_conv1d.py index 0ca1ca421..c39d845f9 100644 --- a/test/registered/layers/mamba/test_causal_conv1d.py +++ b/test/registered/layers/mamba/test_causal_conv1d.py @@ -378,6 +378,62 @@ def test_causal_conv1d_varlen( assert torch.allclose(unpadded_out, out_ref_tensor, rtol=rtol, atol=atol) +def test_causal_conv1d_varlen_mixed_input_and_state_dtype(): + """Initial states may be bf16 even when the current hidden states are fp16.""" + device = get_device() + torch.manual_seed(0) + dim, width = 64, 4 + seqlens = [7, 9] + query_start_loc = torch.tensor([0, 7, 16], dtype=torch.int32, device=device) + x = torch.randn(dim, sum(seqlens), dtype=torch.float16, device=device) + weight = torch.randn(dim, width, dtype=torch.float16, device=device) + bias = torch.randn(dim, dtype=torch.float16, device=device) + + conv_states = torch.randn( + 4, width - 1, dim, dtype=torch.bfloat16, device=device + ).transpose(1, 2) + conv_states_ref = conv_states.clone() + cache_indices = torch.tensor([1, 3], dtype=torch.int32, device=device) + has_initial_state = torch.tensor([True, False], dtype=torch.bool, device=device) + + out = causal_conv1d_fn( + x, + weight, + bias=bias, + conv_states=conv_states, + query_start_loc=query_start_loc, + seq_lens_cpu=torch.tensor(seqlens), + cache_indices=cache_indices, + has_initial_state=has_initial_state, + activation="silu", + ) + + expected = [] + offset = 0 + for i, seqlen in enumerate(seqlens): + state_idx = cache_indices[i] + x_i = x[:, offset : offset + seqlen].unsqueeze(0) + initial_state = ( + conv_states_ref[state_idx].unsqueeze(0).to(x.dtype) + if has_initial_state[i] + else None + ) + out_i, _ = causal_conv1d_ref( + x_i, + weight, + bias, + initial_states=initial_state, + return_final_states=True, + final_states_out=conv_states_ref[state_idx].unsqueeze(0), + activation="silu", + ) + expected.append(out_i.squeeze(0)) + offset += seqlen + + torch.testing.assert_close(out, torch.cat(expected, dim=-1), rtol=1e-2, atol=5e-2) + torch.testing.assert_close(conv_states, conv_states_ref, rtol=1e-2, atol=5e-2) + + if __name__ == "__main__": import sys