"vscode:/vscode.git/clone" did not exist on "d0feea31c70f9540a8993c4e96103a03cd935416"
Unverified Commit 5b55c0be authored by Francesco Fusco's avatar Francesco Fusco Committed by GitHub
Browse files

[Attention] Clarify comment explaining attn_logits +1 dimension (#33427)


Signed-off-by: default avatarFrancesco Fusco <ffu@zurich.ibm.com>
parent 15e0bb9c
......@@ -143,8 +143,8 @@ class TritonMLAImpl(MLACommonImpl[MLACommonMetadata]):
B,
q_num_heads,
num_kv_splits,
# NOTE(lucas) idk why the +1 is here but sglang has it so we
# just mirror that
# NOTE: the +1 stores the LogSumExp (LSE) that the stage2
# kernel uses to merge partial attention outputs across splits.
self.kv_lora_rank + 1,
),
dtype=torch.float32,
......
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