Repository navigation
[JIT][ROCm] Fully warmup GLM-5.3-Flash - #60919
Conversation
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
…vllm into feature/warmup-glm5next
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
simondanielsson
changed the title
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
![claude[bot]](https://avatars.githubusercontent.com/in/1236702?s=60&v=4){kind=small}
|
@simondanielsson #60084 is also open and is approved. Please check that and collaborate. Thanks! |
@GodRishUniverse Thanks! IIUC that fix seems orthogonal to the changes here but indeed it's on the same topic This PR doesn't try to convince triton that the k and v is indeed <2GiB. Instead it adds the option to specialize the kernel on >2GiB buffers during warmup instead of assuming it's always 32-bit addressable |
Cool. Thanks for the quick response! |
| kv_cache = k_cache.kv_cache | ||
| page_size = kv_cache.shape[1] | ||
| buf_numel_per_page = kv_cache.stride(0) | ||
| buf_ptr_range_bytes = triton_ptr_range(kv_cache) |
There was a problem hiding this comment.
Note: This is where the specialization otherwise broke. If not specifying that the buffer it would warmup assuming <2GiB, and at runtime re-compile for >2GiB
|
LGTM. Added comment on the test. |
… bytes Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
| if backend_name != "DEEPSEEK_V32_INDEXER": | ||
| if backend_name not in ( | ||
| "DEEPSEEK_V32_INDEXER", | ||
| "ROCM_AITER_MLA_SPARSE", |
There was a problem hiding this comment.
Can you double check skipping this warmup is okay/needed?
There was a problem hiding this comment.
Skipping is OK but not really related, so I removed it here.
Details if curious:
That specific warmup is only needed for models using Deepseek SWA (dsv4 and dsv41), so I think the actual proper fix is
if backend_name in (ROCM_FLASHMLA_SPARSE_DSV4, ROCM_FLASHMLA_SPARSE_DSV4)
So I'd vouch for removing my change and instead defer this change to another PR as it'd require more validation.
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
simondanielsson
added
the
ready
|
/ci run |
|
✅ Triggered Buildkite CI #98361 for commit |
…next Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
|
/ci run |
|
✅ Triggered Buildkite CI #98429 for commit |
|
Optional, not blocking: the kpool kernels can warm with 2 compiles instead of 8 This PR warms 8 variants of the two kpool kernels: Proposed fix: opt the two kernels out of these specializations: @triton.jit(do_not_specialize_on_alignment=["loc_ptr"])
def _kpool_softmax_rotate_write_cache_kernel(...)
@triton.jit(do_not_specialize=["n_tokens"], do_not_specialize_on_alignment=["tslot_ptr"])
def _kpool_tail_seed_kernel(...)Measured through this PR's warm-up at its head on MI355X: 2 + 6 compiles become 1 + 1, and the warm-up of the two kernels takes 538 ms instead of 769–775 ms per process. In both cases, 24 prefill-write batches after the warm-up compile nothing and write identical cache bytes. With JIT warm-up off ( Setup
|
Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
…next Signed-off-by: simondanielsson <simon.danielsson99@hotmail.com>
|
/ci run |
|
✅ Triggered Buildkite CI #98806 for commit |
dllehr-amd
left a comment
dllehr-amd
left a comment
There was a problem hiding this comment.
Looks good to me. Thanks @simondanielsson!
Overview
It used to JIT several triton kernels during runtime. This PR
Explanation for the triton fix::
kpool_tail_seed_dispatchwe pass the kv cache which can be large. During warmup it specialized on a <2GiB KV cache, whereas at runtime it gets passed a larger KV cache and re-compiles without buffer loads. See https://github.com/ROCm/triton/blob/0ec280cf80dd91e9a86887981a670f2d4541a32b/third_party/amd/backend/compiler.py#L76-L110 for detailsClaims
--jit-monitor-mode errorValidation
Expect lots of runtime compilation on main, none on this branch
VLLM_ROCM_USE_AITER=1 \ VLLM_ROCM_USE_AITER_TRITON_SPARSE_MLA=1 \ VLLM_ROCM_USE_AITER_FUSION_SHARED_EXPERTS=1 \ VLLM_USE_BREAKABLE_CUDAGRAPH=1 \ VLLM_USE_RUST_FRONTEND=1 \ vllm serve zai-org/GLM-5.3-Flash \ --tensor-parallel-size 8 \ --max-num-seqs 512 \ --gpu-memory-utilization 0.8 \ --attention-backend ROCM_AITER_MLA_SPARSE \ --max-num-batched-tokens 16384 \ --jit-monitor-mode warn --jit-monitor-verboseMain: 53 compile keys and runtime compilation
Here we can see the issue: at runtime one of the parameters are specialized as
('*bf16', 'D')(i.e., notDS, so without anSsuffix) which means it's >2GiB. Here's the code that does this.This branch: 167 compile keys and no runtime jit
This branch with
--jit-monitor-mode error: no crash and ok output{"id":"chatcmpl-a5f240f6","object":"chat.completion","created":1792521728,"model":"zai-org/GLM-5.3-Flash","choices":[{"index":0,"message":{"role":"assistant","content":null,"reasoning":"The user wants a short joke and asked me to think for 50 tokens before answering. Let me think of a good short joke.\n\nSome options:\n1. Why don't scientists trust atoms? Because they make up everything!\n2. What do you call a fake noodle? An impasta!\n3. Why did the scarecrow win an award? He was outstanding in his field!\n\nLet me pick one that's classic and reliably funny. The atoms one is a solid choice - it's short"},"logprobs":null,"finish_reason":"length","stop_reason":null,"token_ids":null}],"usage":{"prompt_tokens":34,"total_tokens":134,"completion_tokens":100,"prompt_tokens_details":null,"completion_tokens_details":{"reasoning_tokens":100}},"metrics":null,"system_fingerprint":null,"prompt_logprobs":null,"prompt_token_ids":null,"kv_transfer_params":null,"ec_transfer_params":null}Details
Pull Request Checklist
I used vLLM's
/pr-checklistskill. (Mandatory for agents, optional for humans).AI assistance was used during the creation of this PR.
Design Fit: Minimizes impact on core components, reuses existing functionality, and justifies added complexity.
Testing and Validation: Validates the change and ensures any added tests are meaningful and reliable, with CI coverage or documented CI resource constraints and validation performed outside CI.
Code Quality and Style: Keeps code and comments clear and concise, and updates relevant documentation and examples.
Pull Request Contents: Includes a brief summary and relevant links, supports claims with evidence, explains root causes and implementation trade-offs, and follows the contributing guide.