Execute fused QKVG attention projections - #699
Conversation
|
Codex usage limits have been reached for code reviews. Please check with the admins of this repo to increase the limits by adding credits. |
uuuvn
left a comment
uuuvn
left a comment
There was a problem hiding this comment.
Thank you very much for taking a time to do this!
One thing is that we shouldn't name variables "qkv" when they are in fact "qkvg", other than that looks good, will merge after trymirai/lalamo#340 is merged and a lalamo version is released with it.
| #[uzu_config(super::TokenMixerConfig)] | ||
| pub struct AttentionConfig { | ||
| pub qkv_projection_config: LinearConfig, | ||
| /// Unified QKVG projection; ungated attention omits the G output segment. |
There was a problem hiding this comment.
Ha, I wrote it myself 🤭 Just wanted to emphasize the "G" part may be skipped. But I'll drop this comment in the subsequent update if you feel that way.
| q_dim | ||
| }; | ||
| let projection_dim = if config.has_gate { | ||
| let qkvg_dim = if config.has_gate { |
There was a problem hiding this comment.
Why change "projection" to "qkvg"?
| let key_norm_config = (!is_kv_sharing).then(|| config.key_norm_config.clone()).flatten(); | ||
| let value_norm_config = (!is_kv_sharing).then(|| config.value_norm_config()).flatten(); | ||
| let packed_projection_heads = projection_dim / head_dim; | ||
| let qkvg_heads = qkvg_dim / head_dim; |
There was a problem hiding this comment.
There is no gate heads, qkvg_heads is wrong here
| }; | ||
|
|
||
| pub(super) struct LinearProjection<B: Backend> { | ||
| pub(super) struct QKVGProjection<B: Backend> { |
There was a problem hiding this comment.
Purest misfire on my end, as I'm still getting ramped up with the codebase knowledge. Now fixed.
uuuvn
left a comment
uuuvn
left a comment
There was a problem hiding this comment.
LGTM, will merge when lalamo part merges and releases, thanks
|
In older versions, the kernel computed the offset internally // older qkv_norm.rs kernel
let qkv_stride = total_heads * head_dim;
let offset =
(token_index * qkv_stride)
+ HEAD_START_OFFSET;After these changes, the stride is passed explicitly as self.encode_packed(
key_value,
batch_dim,
0,
2 * self.num_kv_heads, // this is the number of heads. should be the number of elements in the token
encoder,
)should be: - self.encode_packed(key_value, batch_dim, 0, 2 * self.num_kv_heads, encoder)
+ self.encode_packed(key_value, batch_dim, 0, 2 * self.num_kv_heads * self.head_dim, encoder)
The kernel still computes: let offset =
(batch * input_row_stride) // batch * number_of_elements (not number of heads)
+ HEAD_START_OFFSET;
this path only happens when DFlash synchronizes accepted speculative tokens through |
|
Great catch, @gearonixx! Now fixed ✨ |
e68e8a4 to
a4071d1
Compare
8379eb4 to
28f66e9
Compare
3 participants
Consumes Lalamo's fused QKVG projection directly. This eliminates the separate gate-projection dispatch and the gated-input scratch copy while retaining explicit packed-row strides for attention preparation, Q/K/V normalization, and post-attention sigmoid gating.
Benchmarks
Validation
cargo test -p backend-uzu qkv_norm -- --nocapturecargo test -p backend-uzu attention_prepare -- --nocapturecargo test -p backend-uzu sigmoid_gate -- --nocapturecargo clippy -p backend-uzu --lib --tests -- -D warningsDependency
Requires trymirai/lalamo#340 and a Lalamo release containing it.