To increase LX memory usage, torch-spyre implemented hints like
with spyre_hint(tiles={"B": B // b_block_size}):
with spyre_hint(tiles={"H": H // h_block_size}):
with spyre_hint(tiles={"Lq": Lq // q_block_size}):
with spyre_hint(tiles={"Lk": Lk // kv_block_size}):
with spyre_hint(work_div={"H": 4, "Lq": 8, "Lk": 8}):
which should help via steering the working-set right-sizing (or loop rolling/splitting) in the right way.
As part of this issue, we should experiment with these new hints inside the online softmax kernel, evaluate improvements, and investigate potential auto-tuning opportunities.
To increase LX memory usage, torch-spyre implemented hints like
which should help via steering the working-set right-sizing (or loop rolling/splitting) in the right way.
As part of this issue, we should experiment with these new hints inside the online softmax kernel, evaluate improvements, and investigate potential auto-tuning opportunities.