41 lines
1.9 KiB
C++
41 lines
1.9 KiB
C++
/* Copyright 2026 The xLLM Authors. All Rights Reserved.
|
||
|
||
Licensed under the Apache License, Version 2.0 (the "License");
|
||
you may not use this file except in compliance with the License.
|
||
You may obtain a copy of the License at
|
||
|
||
https://github.com/jd-opensource/xllm/blob/main/LICENSE
|
||
|
||
Unless required by applicable law or agreed to in writing, software
|
||
distributed under the License is distributed on an "AS IS" BASIS,
|
||
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||
See the License for the specific language governing permissions and
|
||
limitations under the License.
|
||
==============================================================================*/
|
||
|
||
// Commit: 494f293b5629 · feat · PR #2260 (adapted for Iluvatar BI-V100)
|
||
// gflag definitions for layerwise split KV cache.
|
||
//
|
||
// Usage:
|
||
// --enable_layerwise_split=true --layerwise_split_size=2
|
||
//
|
||
// For Qwen3.5 (Qwen3.6-35B-A3B) on 4× BI-V100 with TP=4:
|
||
// layerwise_split_size=2 means each rank owns ~half the full-attention
|
||
// layers' KV cache, reducing peak per-rank KV memory.
|
||
|
||
#include <gflags/gflags.h>
|
||
|
||
DEFINE_bool(enable_layerwise_split, false,
|
||
"Enable layerwise-split KV cache sharding. When true, each "
|
||
"full-attention layer's KV cache ownership is assigned to a "
|
||
"subset of TP ranks via round-robin. Linear-attention layers "
|
||
"(GDN/DeltaNet) are unaffected — they use conv+temporal state, "
|
||
"not KV cache. Default: false (uniform sharding).");
|
||
|
||
DEFINE_int32(layerwise_split_size, 1,
|
||
"Layer-owner KV cache group size inside each attention TP group. "
|
||
"1 disables layerwise split; values > 1 shard persistent KV by "
|
||
"layer owner and enable layerwise-split communication. The value "
|
||
"must divide attention TP size. Matches upstream xLLM "
|
||
"parallel_config.layerwise_split_size.");
|