[Kernel] [CPU] refactor cpu_attn.py:_run_sdpa_forward for better memory access
#24701
Conversation
Signed-off-by: ignaciosica <mignacio.sica@gmail.com>
ignaciosica
requested review from
WoosukKwon,
robertgshaw2-redhat,
njhill,
ywang96,
comaniac and
alexm-redhat
as code owners
ignaciosica
changed the title
cpu_attn.py: _run_sdpa_forward for better memory access
ignaciosica
changed the title
cpu_attn.py: _run_sdpa_forward for better memory accesscpu_attn.py:_run_sdpa_forward for better memory access
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This is an excellent contribution that significantly improves the performance of CPU attention. The change to reorder the movedim and repeat_interleave operations in _run_sdpa_forward is well-reasoned and clearly explained. By aligning the key/value tensors to a contiguous head-major layout before repeating, you've made the memory access pattern much more efficient for the subsequent scaled_dot_product_attention call, as demonstrated by the impressive performance gain in your benchmarks.
The change is logically sound and I don't see any issues with it. The fact that the outputs match confirms its correctness.
Regarding your question about testing, you've done a great job with local benchmarking. To add a permanent test, you could consider:
-
Unit Test for Correctness: You can add a new test file under the
tests/directory. For an example of a CPU-only test, you can look attests/kernels/attention/test_mla_decode_cpu.py, which uses@pytest.mark.cpu_model. A new test could execute the_run_sdpa_forwardfunction with your changes and compare its output against a reference implementation to ensure numerical correctness for various input shapes. -
Performance Benchmark: To track performance and prevent future regressions, you could investigate adding a benchmark test to the CI. This would formalize the performance gains you've observed.
Overall, this is a high-quality pull request. Great work!
…mory access (vllm-project#24701) Signed-off-by: ignaciosica <mignacio.sica@gmail.com>
…mory access (vllm-project#24701) Signed-off-by: ignaciosica <mignacio.sica@gmail.com>
…mory access (vllm-project#24701) Signed-off-by: ignaciosica <mignacio.sica@gmail.com>
Refactor movement ops in _run_sdpa_forward
This pr moves the
repeat_interleavecall to occur aftermovedim, aligning the key/value tensors with the [B, H, L, D] layout before materializing extra heads. This avoids creating large strided views, which previously forced expensive copies. By repeating along the already head-major axis, we get a contiguous memory layout, producing a friendlier memory layout.In local benchmarks this yields more than 2× faster attention computation with identical results. (tested with different loads scenarios).
I'm new to this project so any feedback or guidance to how to properly unit test this change is greatly appreciated. I've not found many docs on testing/CI.
Profile for offline generation
Setup:
Outputs match!
profile (master) (CPU time total: 45.337s)
Discuss the impact of the movement on society, including the civil rights movement and the civil rights movement in the United States and other countries. Conclude with the significance of the movement in the broader context of social change and equality for all people.
Answer:
Summary of the Industrial Revolution
The Industrial Revolution, which began in Britain in the late 18th century, marked a transformative period in human history. Technologically, it introduced machines, factories, and mass production, significantly increasing productivity and efficiency. Economically, it led to urbanization, the rise of the working class, and the growth of industries such as textiles, steel, and coal. Socially, it reshaped family structures, gender roles, and labor conditions, often leading to exploitation and poor working conditions.
The Industrial Revolution began in Britain due to several factors: the demand for labor in factories, the availability of natural resources, and the need to address growing urbanization. Over the following century, it spread globally, influencing economies, societies, and cultures in unprecedented ways. While the benefits included increased productivity and economic growth, the costs included labor exploitation, environmental damage, and social inequality.
The Theory of Evolution by Natural Selection
The theory of evolution by natural selection is a scientific explanation for the development of life on Earth. It posits that organisms adapt to their environment through genetic variation and natural selection, where individuals with traits that enhance survival and reproduction are more likely to pass on their genes to future generations. Darwin arrived at this theory through his observations of natural phenomena, including the fossil record and the development of species over time. Examples of adaptation in animals and plants include the development of wings in birds and the evolution of antibiotic resistance in bacteria.
The theory of evolution by natural selection has had a profound impact on biology, influencing how we understand life and its diversity. It also reshaped religion and society, as it challenged the idea of a predetermined future and emphasized the importance of natural processes.
Examples of Natural Selection in Action
Two classic examples of natural selection in action are the development of the peppered moth and the evolution of the human eye. The peppered moth's coloration changes in response to pollution, giving it a survival advantage. The human eye's ability to adapt to different light conditions, such as in low-light environments, is an example of natural selection.
Healthy Weekly Meal Plan
To create a healthy weekly meal plan, start with a balanced diet that includes proteins, carbohydrates, and healthy fats. Strategies for grocery shopping on a budget include shopping at local markets and
profile (pr) (CPU time total: 32.625s)
Discuss the impact of the movement on society, including the civil rights movement and the civil rights movement in the United States and other countries. Conclude with the significance of the movement in the broader context of social change and equality for all people.
Answer:
Summary of the Industrial Revolution
The Industrial Revolution, which began in Britain in the late 18th century, marked a transformative period in human history. Technologically, it introduced machines, factories, and mass production, significantly increasing productivity and efficiency. Economically, it led to urbanization, the rise of the working class, and the growth of industries such as textiles, steel, and coal. Socially, it reshaped family structures, gender roles, and labor conditions, often leading to exploitation and poor working conditions.
The Industrial Revolution began in Britain due to several factors: the demand for labor in factories, the availability of natural resources, and the need to address growing urbanization. Over the following century, it spread globally, influencing economies, societies, and cultures in unprecedented ways. While the benefits included increased productivity and economic growth, the costs included labor exploitation, environmental damage, and social inequality.
The Theory of Evolution by Natural Selection
The theory of evolution by natural selection is a scientific explanation for the development of life on Earth. It posits that organisms adapt to their environment through genetic variation and natural selection, where individuals with traits that enhance survival and reproduction are more likely to pass on their genes to future generations. Darwin arrived at this theory through his observations of natural phenomena, including the fossil record and the development of species over time. Examples of adaptation in animals and plants include the development of wings in birds and the evolution of antibiotic resistance in bacteria.
The theory of evolution by natural selection has had a profound impact on biology, influencing how we understand life and its diversity. It also reshaped religion and society, as it challenged the idea of a predetermined future and emphasized the importance of natural processes.
Examples of Natural Selection in Action
Two classic examples of natural selection in action are the development of the peppered moth and the evolution of the human eye. The peppered moth's coloration changes in response to pollution, giving it a survival advantage. The human eye's ability to adapt to different light conditions, such as in low-light environments, is an example of natural selection.
Healthy Weekly Meal Plan
To create a healthy weekly meal plan, start with a balanced diet that includes proteins, carbohydrates, and healthy fats. Strategies for grocery shopping on a budget include shopping at local markets and