Skip to content

🤔 experiment with torch-spyre WSRS hints and integrate them into online softmax #329

Description

@bringlein

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.

Metadata

Metadata

Labels

vllm-spyre-attention-backendvllm-spyre-nextRelated to enabling the next-generation of the Spyre stack built on top of torch-spyre

Type

No type

Fields

No fields configured for issues without a type.

Projects

Status
Ready

Milestone

No milestone

Relationships

None yet

Development

No branches or pull requests

Issue actions