From 3101c1258c4d5dc4884868259278fb6370a92617 Mon Sep 17 00:00:00 2001 From: Lewis <63569348+TTThanos@users.noreply.github.com> Date: Thu, 16 Jul 2026 06:06:55 +0800 Subject: [PATCH] [DSv4] Use BF16 instead of FP32 for indexer score computation (#30012) MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: 百麒 --- python/sglang/srt/layers/attention/dsv4/indexer.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/python/sglang/srt/layers/attention/dsv4/indexer.py b/python/sglang/srt/layers/attention/dsv4/indexer.py index e11fbdeec..4fe15baf8 100644 --- a/python/sglang/srt/layers/attention/dsv4/indexer.py +++ b/python/sglang/srt/layers/attention/dsv4/indexer.py @@ -87,14 +87,14 @@ def fp8_paged_mqa_logits_torch( kv_values_raw = kvcache_gathered[..., :SCALE_OFFSET].contiguous() kv_values_fp8 = kv_values_raw.view(dtype=FP8_DTYPE) - kv_values = kv_values_fp8.to(torch.float32) + kv_values = kv_values_fp8.to(torch.bfloat16) kv_values = kv_values.reshape(batch_size, max_num_pages * block_size, head_dim) kv_scales_raw = kvcache_gathered[..., SCALE_OFFSET:].contiguous() kv_scales = kv_scales_raw.view(dtype=torch.float32) kv_scales = kv_scales.reshape(batch_size, max_num_pages * block_size) - q_float = q_fp8[:, 0].to(torch.float32) + q_float = q_fp8[:, 0].to(torch.bfloat16) scores = torch.bmm(kv_values, q_float.transpose(1, 2)) scores = F.relu(scores) scores = scores * weight.unsqueeze(1) @@ -199,13 +199,13 @@ def fp8_paged_mqa_logits_torch_sm120( kv_value_raw = kvcache_gathered[..., :SCALE_OFFSET] kv_scale_raw = kvcache_gathered[..., SCALE_OFFSET:] - kv_value = kv_value_raw.contiguous().view(dtype=FP8_DTYPE).to(torch.float32) + kv_value = kv_value_raw.contiguous().view(dtype=FP8_DTYPE).to(torch.bfloat16) kv_value = kv_value.view(batch_size, max_padded_seq, head_dim) kv_scale = kv_scale_raw.contiguous().view(dtype=torch.float32) kv_scale = kv_scale.view(batch_size, max_padded_seq) - q = q_fp8[:, 0].to(torch.float32) + q = q_fp8[:, 0].to(torch.bfloat16) score = torch.bmm(kv_value, q.transpose(1, 2))