From a19f05180094563bd02a23ad9c1aba768535deba Mon Sep 17 00:00:00 2001 From: Daniel Maddern Date: Wed, 19 Aug 2026 22:10:39 +0700 Subject: [PATCH] Cast qkv to common dtype before SDPA (q/k/v match) --- src/h3_blackwell_runtime/qwen3vl_vision.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/h3_blackwell_runtime/qwen3vl_vision.py b/src/h3_blackwell_runtime/qwen3vl_vision.py index dd28104..7ed230a 100644 --- a/src/h3_blackwell_runtime/qwen3vl_vision.py +++ b/src/h3_blackwell_runtime/qwen3vl_vision.py @@ -265,10 +265,16 @@ class _VisionAttention(nn.Module): def forward(self, x: torch.Tensor, cu_seqlens: torch.Tensor, position_embeddings: torch.Tensor) -> torch.Tensor: seq_length = x.shape[0] + # Cast qkv to a common dtype (the weights may be bf16 or fp32). + qkv_dtype = self.qkv_weight.dtype + qkv = F.linear(x.to(qkv_dtype), self.qkv_weight, self.qkv_bias) query_states, key_states, value_states = ( - F.linear(x, self.qkv_weight, self.qkv_bias).reshape(seq_length, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0) + qkv.reshape(seq_length, 3, self.num_heads, self.head_dim).permute(1, 0, 2, 3).unbind(0) ) - query_states, key_states = _apply_rope_vision(query_states, key_states, position_embeddings) + # RoPE in fp32 to avoid bf16 precision loss, then cast back. + orig_dtype = query_states.dtype + query_states, key_states = _apply_rope_vision(query_states.float(), key_states.float(), position_embeddings.float()) + query_states, key_states = query_states.to(orig_dtype), key_states.to(orig_dtype) lengths = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist() attn_outputs = [] for q, k, v in zip(