Conversation
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
| ((total_seq_lens, head_dim), fp8_dtype), | ||
| ((total_seq_lens, 4), torch.uint8), | ||
| ((total_seq_lens, head_dim), FP8_DTYPE), | ||
| ((total_seq_lens, 1), torch.float32), |
There was a problem hiding this comment.
Is there any context on this? (both for the 4 -> 1 modification and uint8 -> fp32)
AndreasKaratzas
added
the
ready
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
Signed-off-by: tjtanaa <tunjian.tan@embeddedllm.com>
| _ON_GFX942 = False | ||
| _ON_GFX950 = False | ||
|
|
||
| FP8_DTYPE = current_platform.fp8_dtype() |
There was a problem hiding this comment.
I have pre-evaluate the platform DTYPE as the current_platform.fp8_dtype() is a host synchronize op. We should try to invoke it as fewer as possible.
|
For coordination: draft #52882 is intended as a more complete superseding |
current_platform.fp8_dtype() cannot change after import, but the sparse MLA indexer ops call it on every invocation - once per DSA layer per forward, including inside the fused QK prologue custom op. vllm-project#46172 added the module-level FP8_DTYPE constant but left the call sites untouched. Route every use in rocm_aiter_mla_sparse.py through FP8_DTYPE and drop the per-call locals, so exactly one lookup survives, at import.
|
This pull request has merge conflicts that must be resolved before it can be |
current_platform.fp8_dtype() cannot change after import, but the sparse MLA indexer ops call it on every invocation - once per DSA layer per forward. vllm-project#46172 added the module-level FP8_DTYPE constant but left the call sites untouched. Route every use in rocm_aiter_mla_sparse.py through FP8_DTYPE and drop the per-call locals, so exactly one lookup survives, at import. The fused QK prologue's two call sites are not touched here: vllm#51315 moves that impl and its fake into _aiter_ops.py, where they already read that file's FP8_DTYPE constant. Signed-off-by: Sriram Kumar <sriramkumar.kishorekumar@amd.com>
3 participants
Purpose
This is validated using aiter
v0.1.16.post2which is the version used in thevllm/vllm-openai-rocmOptimize the sparse indexer for topk-token = 2048 following speed of light implementation from ATOM https://github.com/ROCm/ATOM/blob/main/atom/plugin/vllm/attention/layer_sparse_mla.py, e.g. DSv3.2 model, with aiter's
cp_gather_indexer_k_quant_cacheandindexer_k_quant_and_cacheandtop_k_per_row_decode.This alignment is to capitalize the long-term benefits from AITER as ATOM is invested in optimizing these AITER kernels.
Three small kernels are intended to be upstream together due to its small changes, and the long-term benefits brought these kernels. All three kernels integration also provides performance gain.
Test Plan
lm-eval of DSv3.2
Test Result
Block size 1
Block size 64
lm-eval of DeepSeek V4 (represent topk = 1024)
DeepSeek V3.2
DeepSeek V4-Flash
DeepSeek V4-Pro
Essential Elements of an Effective PR Description Checklist
supported_models.mdandexamplesfor a new model.