daniel/h3-blackwell-runtime
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Cast qkv to common dtype before SDPA (q/k/v match)

Browse source
This commit is contained in:
Daniel Maddern 2026-08-19 22:10:39 +07:00
parent efc5fd9bc1
commit a19f051800
1 changed files with 8 additions and 2 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

10
src/h3_blackwell_runtime/qwen3vl_vision.py
View file

@ -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(