Summary
_compute_dq_write_order_from_block_mask overhead in compiled create_block_mask
Measured end-to-end: torch.compile(create_block_mask)(..., compute_dq_write_order=True) vs without.
| B | H | S | create_block_mask (ms) | w/ dq_write_order (ms) | delta (ms) | overhead % |
|---|---|---|---|---|---|---|
| 1 | 8 | 2048 | 0.066 | 0.077 | +0.012 | +17.7% |
| 1 | 8 | 8192 | 0.489 | 0.495 | +0.006 | +1.3% |
| 1 | 8 | 16384 | 2.136 | 2.149 | +0.014 | +0.6% |
| 1 | 8 | 32768 | 8.222 | 8.241 | +0.019 | +0.2% |
| 1 | 16 | 8192 | 0.846 | 0.856 | +0.010 | +1.2% |
| 1 | 16 | 16384 | 4.139 | 4.154 | +0.015 | +0.4% |
| 1 | 16 | 32768 | 16.301 | 16.323 | +0.023 | +0.1% |
| 4 | 8 | 8192 | 2.162 | 2.169 | +0.007 | +0.3% |
| 4 | 8 | 32768 | 32.459 | 32.488 | +0.030 | +0.1% |
| 4 | 16 | 8192 | 4.190 | 4.204 | +0.013 | +0.3% |
| 4 | 16 | 16384 | 16.194 | 16.216 | +0.022 | +0.1% |
| 4 | 16 | 32768 | 64.794 | 64.859 | +0.066 | +0.1% |
| 8 | 8 | 8192 | 4.191 | 4.203 | +0.013 | +0.3% |
| 8 | 8 | 32768 | 64.800 | 64.868 | +0.068 | +0.1% |
| 8 | 16 | 8192 | 8.246 | 8.263 | +0.017 | +0.2% |
| 8 | 16 | 16384 | 32.256 | 32.283 | +0.027 | +0.1% |
| 8 | 16 | 32768 | 129.485 | 129.602 | +0.117 | +0.1% |
Mask types tested: causal, sliding_window, doc_mask_mod (numbers above are for causal; other masks show identical overhead).
Summary: At S >= 8192 the overhead is consistently < 0.3%
Looked at generated kernels we are able to fuse pieces with exist q_block calc and we have 1 more kernel for the exclusive prefix sum which is
Pull Request resolved: #174813
Approved by: https://github.com/liangel-02
SOCIAL SHARE CARD GENERATOR