[XPU] Support softmax_lse in sgl_kernel::fwd API (#33840)

Co-authored-by: Ma Mingfei <mingfei.ma@intel.com>
This commit is contained in:
Xuan Liao
2026-08-24 12:05:15 +08:00
committed by GitHub
co-authored by Ma Mingfei
parent 7a7b655ddf
commit acba8921bf
@@ -73,21 +73,13 @@ def register_fake_ops():
num_kv_splits,
pack_gqa,
sm_margin,
out=None,
out,
softmax_lse,
):
total_q = q.shape[0]
num_heads_q = q.shape[1]
head_size_v = v.shape[-1]
if out is None:
out = q.new_empty(total_q, num_heads_q, head_size_v)
softmax_lse = q.new_empty(num_heads_q, total_q, dtype=torch.float32)
# out_accum and softmax_lse_accum are intermediate split-kv buffers;
# they are only read when num_kv_splits > 1, which is determined at
# runtime. Return empty tensors with correct rank so downstream ops
# that index into the list do not fail shape propagation.
out_accum = q.new_empty(0)
softmax_lse_accum = q.new_empty(0, dtype=torch.float32)
return (out, softmax_lse, out_accum, softmax_lse_accum)
# The op writes into the caller-provided ``out`` and ``softmax_lse``
# buffers in place and returns nothing. ``softmax_lse`` is ``None`` when
# the caller does not request the logsumexp.
return None
@torch.library.register_fake("sgl_kernel::flash_mla_decode")
def _(