Add chunking strategy for fp8_paged_mqa_logits - #398
Conversation
It can probably resolve the OOM issue of DeepSeek. Jianan is going to have a try. |
There was a problem hiding this comment.
Pull request overview
Adds configurable batch chunking to the Xe20 FP8 paged MQA logits GEMM path.
Changes:
- Adds a configurable 512 MiB default chunk budget.
- Processes batches with chunk-local intermediates.
- Adds optional verbose chunk diagnostics.
💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.
| " chunks)"); | ||
| } | ||
|
|
||
| for (int start = 0; start < B_next; start += chunk_b) { |
There was a problem hiding this comment.
avoid launching kernels in the loop, fuse into a larger kernel instead
There was a problem hiding this comment.
Thanks for the suggestion. My agent says the loop introduces limited kernel launch overhead (~2% in the DeepSeek V4 case). I have added a TODO here to indicate an opportunity for optimizations. I think we can keep the current implementation to resolve the OOM issue which is blocking.
There was a problem hiding this comment.
After an offline discussion, we decide to improve the kernel implementation for better performance so I have converted this PR to draft and will work on it later. Thanks.
export SGL_KERNEL_FP8_PAGED_MQA_CHUNK_MB=512(default is 512)