Unverified Commit 75fb112d authored by Wentao Ye's avatar Wentao Ye Committed by GitHub
Browse files

[Bug] Fix `returned_lse` not Defined issue (#25106)


Signed-off-by: default avataryewentao256 <zhyanwentao@126.com>
Co-authored-by: default avatarTyler Michael Smith <tyler@neuralmagic.com>
parent 38db529f
...@@ -206,12 +206,11 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]): ...@@ -206,12 +206,11 @@ class CutlassMLAImpl(MLACommonImpl[MLACommonMetadata]):
) )
if H < MAX_HEADS: if H < MAX_HEADS:
# Extract the subsets of the outputs
returned_lse = lse[:, :H].contiguous(
) if self.need_to_return_lse_for_decode else lse
out = out[:, :H] out = out[:, :H]
if self.need_to_return_lse_for_decode:
lse = lse[:, :H].contiguous()
return out, returned_lse return out, lse
def _forward_decode( def _forward_decode(
self, self,
......
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment