Skip to content

Commit 09c76d0

Browse files
authored
docs(functions): disambiguate computed attn_bias; co-locate PackedMHA test
Rename the computed attention bias in the module docstring pseudo-code to attn_bias so it is not confused with the new formal `bias` input, and move the fallback-function test next to its implementation under src/mobius/functions/ to match the existing convention.
1 parent 864c04f commit 09c76d0

2 files changed

Lines changed: 5 additions & 4 deletions

File tree

src/mobius/functions/packed_multi_head_attention.py

Lines changed: 5 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -17,16 +17,17 @@
1717
1818
# Build block-diagonal bias (0 for same segment, -inf for different)
1919
same = Equal(segment_ids[:, None], segment_ids[None, :])
20-
bias = Where(same, 0.0, -10000.0)
20+
attn_bias = Where(same, 0.0, -10000.0)
2121
2222
# Standard Attention
23-
output = Attention(query, key, value, bias,
23+
output = Attention(query, key, value, attn_bias,
2424
q_num_heads=<num_heads>, kv_num_heads=<num_heads>,
2525
scale=<scale>)
2626
2727
The optional ``bias`` (slot 4) and ``token_offset`` inputs are consumed by
28-
the native kernel but are unused by the fallback body (segment boundaries
29-
from ``cumulative_sequence_length`` are sufficient). They are still declared
28+
the native kernel but are unused by the fallback body (the ``attn_bias``
29+
computed above from ``cumulative_sequence_length`` is sufficient, and is
30+
unrelated to the formal ``bias`` input). They are still declared
3031
as formal inputs to preserve the positional slot alignment expected by the
3132
ORT ``PackedMultiHeadAttention`` signature.
3233
File renamed without changes.

0 commit comments

Comments
 (0)