Emit a single SliceUpdate for linear (non-ring) cache.
(self, P: "MLXProgramBuilder", cache_slot, update_slot, start_slot)
| 235 | return cache_slot |
| 236 | |
| 237 | def _emit_linear(self, P: "MLXProgramBuilder", cache_slot, update_slot, start_slot): |
| 238 | """Emit a single SliceUpdate for linear (non-ring) cache.""" |
| 239 | update_meta = self.update.meta.get("val") |
| 240 | stop_slot = emit_stop_position( |
| 241 | P, |
| 242 | start=start_slot, |
| 243 | length_tensor=update_slot, |
| 244 | length_dim=2, # S_step is dim 2 in [B, H, S_step, D] |
| 245 | length_meta=update_meta, |
| 246 | ) |
| 247 | |
| 248 | # This updates cache[:, :, start:stop, :] = update |
| 249 | # SliceUpdateNode on axis=2 |
| 250 | # cache is [B, H, S, D], update is [B, H, S_step, D] |
| 251 | P.emit( |
| 252 | SliceUpdateNode( |
| 253 | dst=P.slot_to_tid(cache_slot), |
| 254 | update=P.slot_to_tid(update_slot), |
| 255 | out=P.slot_to_tid(cache_slot), |
| 256 | axis=IntOrVid.from_literal(2), # S dimension in [B, H, S, D] |
| 257 | start=P.to_int_or_vid(start_slot), |
| 258 | stop=P.to_int_or_vid(stop_slot), |
| 259 | ) |
| 260 | ) |
| 261 | |
| 262 | def _emit_ring_buffer( |
| 263 | self, P: "MLXProgramBuilder", cache_slot, update_slot, start_slot |
no test coverage detected