[HiSparse] Add MHA hisparse support for MiniMax M3 (#31446)

Co-authored-by: Guangda Liu <bingps@users.noreply.github.com>
This commit is contained in:
Guangda Liu
2026-09-22 13:28:03 +08:00
committed by GitHub
co-authored by Guangda Liu
parent 095e45100b
commit 04c0913434
27 changed files with 1016 additions and 118 deletions
@@ -129,6 +129,7 @@ The NVIDIA Blackwell recipes are validated single-node: **B200 at `--tp 8`** and
- **Memory**: `--mem-fraction-static` reserves GPU memory for weights + KV pool; the rest is prefill **activation headroom**. The value scales with *free* memory per GPU (card capacity minus per-GPU weight), so it tracks the card more than the TP degree: **`0.65` on B200** (180 GB — less headroom once weights are resident) and **`0.75` on the larger-memory B300 / GB300** (`0.80` on AMD). Lower TP packs more weight per GPU, so a tighter config needs a *lower* value — B200 needs `0.65` even at `--tp 4`. Raising it past the validated value is fine only for low-concurrency single-stream serving; it OOMs under high concurrency or long context.
- **Long context (32K+)**: keep `--mem-fraction-static` at the platform default and raise `--chunked-prefill-size` to `16384`. Decode TPOT stays roughly flat in context length thanks to sparse attention; 1K128K prompts are validated.
- **HiSparse for decode capacity**: on NVIDIA CUDA, HiSparse keeps the three dense layers on GPU, moves the 57 sparse-layer K/V caches to pinned host memory, and feeds selected block IDs directly to the swap-in kernel. For the released four-KV-head model, use `--tp 4` or greater, `--disable-radix-cache`, and `device_buffer_size >= 2048`. Enable it with `--enable-hisparse --hisparse-config='{"device_buffer_size":4096,"host_to_device_ratio":2}'` on the Triton launch command.
- **Scaling TP**: B200 is documented at `--tp 8`; B300 / GB200 / GB300 at `--tp 4` (the single-node cross-family common denominator). On an 8-GPU B300 host you can also raise to `--tp 8` for more throughput / KV headroom.
- **Expert parallelism**: to trade latency for throughput add `--ep` (see [Expert Parallelism Deployment](../../../docs/advanced_features/expert_parallelism)). On AMD, set `--ep` equal to `--tp`. Shared-experts fusion is automatically disabled when EP > 1; on AMD standard EP the server also disables `--enable-aiter-allreduce-fusion` automatically to preserve accuracy.
- `--trust-remote-code` is required to load the MiniMax config / processor classes.
+10 -1
View File
@@ -6,7 +6,7 @@ metatags:
HiSparse reduces per-request GPU memory consumption during the decode phase by maintaining only a small "hot" KV buffer on GPU while keeping complete KV data in CPU pinned memory. Combined with PD disaggregation, it enables significantly higher decode concurrency.
> **Prerequisites**: HiSparse works with models that use **DeepSeek Sparse Attention (DSA)** architectures (e.g., DeepSeek-V3.2, GLM-5.1) and **DeepSeek V4**. These models natively select a subset of tokens for attention, making it possible to keep only the top-k KV on GPU while storing the full KV in host memory — without accuracy loss. Additionally, HiSparse currently requires **PD disaggregation mode** and is enabled on the **decode instance** only.
> **Prerequisites**: HiSparse works with models that use **DeepSeek Sparse Attention (DSA)** architectures (e.g., DeepSeek-V3.2, GLM-5.1), **DeepSeek V4**, and **MiniMax M3**. These models natively select a subset of tokens for attention, making it possible to keep only the top-k KV on GPU while storing the full KV in host memory — without accuracy loss. Additionally, HiSparse currently requires **PD disaggregation mode** and is enabled on the **decode instance** only.
## Why HiSparse?
@@ -165,6 +165,15 @@ python3 -m sglang.launch_server \
> **Note**: For DSA models, `--kv-cache-dtype` defaults to `auto`, which resolves to `fp8_e4m3` on SM100+ (Blackwell) and `bfloat16` on older architectures. The DSA decode backend is automatically selected based on KV dtype (`bfloat16` → `flashmla_sparse`, `fp8_e4m3` → `flashmla_kv`), except for GLM DSA models on SM120/SM121 with `fp8_e4m3`, which use `flashinfer_sparse_mla`. DSA backend flags apply only to DSA models; DeepSeek V4 uses its own `dsv4` attention backend.
### MiniMax M3
Dense-layer K/V and index K stay on GPU; sparse-layer K/V use host memory plus a GPU working set.
- Use TP 4 or greater, with the same TP size and PP 1 on both PD instances.
- Use `--attention-backend triton`, `--mm-attention-backend triton_attn`, `--disable-prefill-cuda-graph`, and `--disable-radix-cache`.
- Set `device_buffer_size` to at least 2048 in `--hisparse-config`; `top_k` does not override the model's selection width.
- PD retraction backup is not supported. Use `--num-reserved-decode-tokens` to reserve capacity for the expected output length.
### Benchmark
```bash Command
@@ -303,13 +303,26 @@ template <int NUM_TOP_K, int HOT_BUFFER_SIZE>
struct SmemLayout {
static constexpr int HASH_SIZE = NUM_TOP_K * 2;
static constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + WARP_SIZE - 1) / WARP_SIZE;
// int32_t region: top_k_tokens + chunk_offset + evict_chunk_offset + hash_keys + total_hits + newest_hit
// int32_t region: top_k_tokens + chunk offsets + hash keys + hit counters
static constexpr int TOTAL_INT32 = NUM_TOP_K + (NUM_BUFFER_CHUNKS + 1) + (NUM_BUFFER_CHUNKS + 1) + HASH_SIZE + 2;
// int16_t region: lru_slots_out + hash_vals
static constexpr int TOTAL_INT16 = HOT_BUFFER_SIZE + HASH_SIZE;
static constexpr size_t BYTES = TOTAL_INT32 * sizeof(int32_t) + TOTAL_INT16 * sizeof(int16_t);
};
template <int SPARSE_BLOCK_SIZE, bool TopKIsBlocks>
__device__ __forceinline__ int32_t resolve_selected_token(const int32_t* top_k, int32_t token_index) {
if constexpr (TopKIsBlocks) {
const int32_t block_index = top_k[token_index / SPARSE_BLOCK_SIZE];
if (block_index < 0) {
return -1;
}
return block_index * SPARSE_BLOCK_SIZE + token_index % SPARSE_BLOCK_SIZE;
} else {
return top_k[token_index];
}
}
// Each block processes one request
// req_pool_indices and seq_lens can each be int32_t or int64_t
// Layout: [HOT_BUFFER_SIZE slots for LRU] + [page_size slots for newest token]
@@ -319,23 +332,28 @@ struct SmemLayout {
// false -> generic byte-stride: device + host both linear, stride = item_size_bytes
// true -> DSv4 page-padded device + page-padded host (kvcacheio.cuh constants)
//
// TopKIsBlocks makes the kernel consume block ids directly. It resolves token
// positions in registers and writes the flattened token-slot table expected by
// sparse attention without materializing an intermediate token-index tensor.
// RecordMissPlan records this step's miss plan (miss_src/dst = host/device loc
// per miss, miss_count per request) for shared-index skip layers to replay via
// copy_cache_planned_kernel. SkipIO elides only the KV byte movement (timing
// probe; output is garbage). Both are compile-time flags so the production
// (false, false) instantiation stays byte-identical.
// probe; output is garbage). These are compile-time flags, so inactive paths
// are removed from each specialization.
template <
int BLOCK_SIZE,
int NUM_TOP_K,
int HOT_BUFFER_SIZE,
bool IsMLA,
bool IsDsv4Layout,
int SPARSE_BLOCK_SIZE,
bool TopKIsBlocks,
bool RecordMissPlan,
bool SkipIO,
typename SeqLensT,
typename ReqPoolIndicesT>
__global__ void load_cache_to_device_buffer_kernel(
const int32_t* __restrict__ top_k_tokens,
const int32_t* __restrict__ top_k,
int32_t* __restrict__ device_buffer_tokens,
const int64_t* __restrict__ host_cache_locs,
const int32_t* __restrict__ device_buffer_locs,
@@ -351,7 +369,7 @@ __global__ void load_cache_to_device_buffer_kernel(
int64_t buffer_stride_0,
int64_t host_stride,
int64_t lru_slot_stride_0,
int64_t top_k_tokens_stride,
int64_t top_k_stride,
int64_t top_k_device_locs_stride,
int64_t page_size,
int64_t item_size_bytes,
@@ -360,9 +378,12 @@ __global__ void load_cache_to_device_buffer_kernel(
int32_t* __restrict__ miss_count_out,
int64_t plan_stride) {
static_assert(!IsDsv4Layout || IsMLA, "DSv4 page-padded layout is K-only (MLA).");
// todo hisparse: support page wise sparsity
static_assert(SPARSE_BLOCK_SIZE > 0, "SPARSE_BLOCK_SIZE must be positive.");
// Cache residency and LRU replacement remain token-granular even when the
// sparse-attention selection arrives as block ids.
constexpr int NUM_TOP_K_TOKENS = NUM_TOP_K * (TopKIsBlocks ? SPARSE_BLOCK_SIZE : 1);
constexpr int NUM_WARPS = BLOCK_SIZE / WARP_SIZE;
constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K + WARP_SIZE - 1) / WARP_SIZE;
constexpr int NUM_TOKEN_CHUNKS = (NUM_TOP_K_TOKENS + WARP_SIZE - 1) / WARP_SIZE;
constexpr int NUM_BUFFER_CHUNKS = (HOT_BUFFER_SIZE + WARP_SIZE - 1) / WARP_SIZE;
const int bid = blockIdx.x;
@@ -372,7 +393,7 @@ __global__ void load_cache_to_device_buffer_kernel(
// CUDA graph pads the batch to a captured size. Keep padded output rows
// invalid without a separate fill kernel.
if (bid >= num_real_reqs[0]) {
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
req_top_k_device_locs[i] = -1;
}
return;
@@ -386,7 +407,7 @@ __global__ void load_cache_to_device_buffer_kernel(
const int64_t seq_len = seq_lens[bid];
// Calculate offsets for this request
const int32_t* req_top_k_tokens = top_k_tokens + bid * top_k_tokens_stride;
const int32_t* req_top_k = top_k + bid * top_k_stride;
const int64_t buffer_offset = rid * buffer_stride_0;
int32_t* req_device_buffer_tokens = device_buffer_tokens + buffer_offset;
@@ -396,14 +417,16 @@ __global__ void load_cache_to_device_buffer_kernel(
// Fast path: short sequences have all tokens in the device buffer in order.
if (seq_len <= HOT_BUFFER_SIZE) {
const int count = (seq_len < NUM_TOP_K) ? static_cast<int>(seq_len) : NUM_TOP_K;
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
const int count = (seq_len < NUM_TOP_K_TOKENS) ? static_cast<int>(seq_len) : NUM_TOP_K_TOKENS;
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
int32_t device_loc = -1;
if (i < count) {
int32_t token_pos = req_top_k_tokens[i];
if (token_pos >= 0) {
const int32_t token_pos = resolve_selected_token<SPARSE_BLOCK_SIZE, TopKIsBlocks>(req_top_k, i);
if constexpr (TopKIsBlocks) {
if (token_pos >= 0 && token_pos < seq_len) {
device_loc = req_device_buffer_locs[token_pos];
}
} else if (i < count && token_pos >= 0) {
device_loc = req_device_buffer_locs[token_pos];
}
req_top_k_device_locs[i] = device_loc;
}
@@ -418,21 +441,21 @@ __global__ void load_cache_to_device_buffer_kernel(
// Dynamic shared memory layout: int32_t arrays first, then int16_t arrays.
extern __shared__ char smem_raw[];
using Layout = SmemLayout<NUM_TOP_K, HOT_BUFFER_SIZE>;
using Layout = SmemLayout<NUM_TOP_K_TOKENS, HOT_BUFFER_SIZE>;
constexpr int HASH_SIZE = Layout::HASH_SIZE;
int32_t* smem_i32 = reinterpret_cast<int32_t*>(smem_raw);
// Top-k token positions; reused as miss-token scratch in the copy phase
int32_t* s_top_k_tokens = smem_i32;
// Prefix-sum offsets for hit counting and miss counting
int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K;
int32_t* s_chunk_offset = s_top_k_tokens + NUM_TOP_K_TOKENS;
// Prefix-sum offsets for evictable counting
int32_t* s_evict_chunk_offset = s_chunk_offset + (NUM_BUFFER_CHUNKS + 1);
// Open-addressing hash table: top-k token_id -> top-k index (keys)
int32_t* s_hash_keys = s_evict_chunk_offset + (NUM_BUFFER_CHUNKS + 1);
// Scalar counters
int32_t& s_total_hits = s_hash_keys[HASH_SIZE];
int32_t& s_newest_hit = s_hash_keys[HASH_SIZE + 1];
int32_t& s_total_misses = s_hash_keys[HASH_SIZE + 1];
int16_t* smem_i16 = reinterpret_cast<int16_t*>(smem_i32 + Layout::TOTAL_INT32);
// Compacted slot ordering: [hits fwd-> ... <-evictables bwd]
@@ -443,7 +466,7 @@ __global__ void load_cache_to_device_buffer_kernel(
// Initialize shared memory: counters, hash table, prefix-sum offsets.
if (tid == 0) {
s_total_hits = 0;
s_newest_hit = 0;
s_total_misses = 0;
}
for (int i = tid; i < HASH_SIZE; i += BLOCK_SIZE) {
s_hash_keys[i] = HASH_EMPTY;
@@ -458,14 +481,20 @@ __global__ void load_cache_to_device_buffer_kernel(
const int32_t newest_token = seq_len - 1;
// Insert top-k tokens into shared-memory hash table.
for (int i = tid; i < NUM_TOP_K; i += BLOCK_SIZE) {
int32_t token_idx = req_top_k_tokens[i];
for (int i = tid; i < NUM_TOP_K_TOKENS; i += BLOCK_SIZE) {
const int32_t token_idx = resolve_selected_token<SPARSE_BLOCK_SIZE, TopKIsBlocks>(req_top_k, i);
if constexpr (TopKIsBlocks) {
if (token_idx < 0 || token_idx >= seq_len) {
s_top_k_tokens[i] = TOKEN_HIT;
req_top_k_device_locs[i] = -1;
continue;
}
}
if (token_idx == newest_token) {
// If topk includes the latest token, bind its canonical occurrence to newest_slot (at HOT_BUFFER_SIZE) and mark
// it as a hit. newest_slot is at the first position of the extra page, excluded from LRU tracking.
s_top_k_tokens[i] = TOKEN_HIT;
req_top_k_device_locs[i] = req_device_buffer_locs[newest_slot];
s_newest_hit = 1;
} else {
int slot = hash_slot(token_idx, HASH_SIZE);
while (true) {
@@ -580,7 +609,7 @@ __global__ void load_cache_to_device_buffer_kernel(
const int chunk_token_start = chunk_idx * WARP_SIZE;
const int my_token_idx = chunk_token_start + lane_id;
const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K);
const bool has_valid_token = has_valid_chunk && (my_token_idx < NUM_TOP_K_TOKENS);
int32_t my_token = 0;
bool is_miss = false;
@@ -611,6 +640,9 @@ __global__ void load_cache_to_device_buffer_kernel(
#else
total_misses = warp_inclusive_scan(s_chunk_offset, lane_id, chunk_idx + 1, NUM_TOKEN_CHUNKS + 1, total_misses);
#endif
if (tid == 0) {
s_total_misses = total_misses;
}
}
__syncthreads();
@@ -632,7 +664,7 @@ __global__ void load_cache_to_device_buffer_kernel(
}
__syncthreads();
total_misses = NUM_TOP_K - s_total_hits - s_newest_hit;
total_misses = s_total_misses;
if constexpr (RecordMissPlan) {
if (tid == 0) {
miss_count_out[bid] = total_misses;
@@ -695,10 +727,12 @@ template <
int HOT_BUFFER_SIZE,
bool IsMLA,
bool IsDsv4Layout,
int SPARSE_BLOCK_SIZE,
bool TopKIsBlocks,
bool RecordMissPlan,
bool SkipIO>
void load_cache_to_device_buffer(
tvm::ffi::TensorView top_k_tokens,
tvm::ffi::TensorView top_k,
tvm::ffi::TensorView device_buffer_tokens,
tvm::ffi::TensorView host_cache_locs,
tvm::ffi::TensorView device_buffer_locs,
@@ -718,7 +752,8 @@ void load_cache_to_device_buffer(
tvm::ffi::TensorView miss_count_out) {
using namespace host;
const int64_t bs = top_k_tokens.shape()[0];
constexpr int NUM_TOP_K_TOKENS = NUM_TOP_K * (TopKIsBlocks ? SPARSE_BLOCK_SIZE : 1);
const int64_t bs = top_k.shape()[0];
const int64_t host_stride = host_cache_locs.shape()[1];
// Miss-plan side outputs; 0-dim sentinels when RecordMissPlan is false.
int64_t* const miss_src_ptr = RecordMissPlan ? static_cast<int64_t*>(miss_src_out.data_ptr()) : nullptr;
@@ -730,9 +765,9 @@ void load_cache_to_device_buffer(
}
const int64_t buffer_stride_0 = device_buffer_tokens.strides()[0];
const int64_t lru_slot_stride_0 = lru_slots.strides()[0];
const int64_t top_k_tokens_stride = top_k_tokens.strides()[0];
const int64_t top_k_stride = top_k.strides()[0];
const int64_t top_k_device_locs_stride = top_k_device_locs.strides()[0];
const auto kernel_device = top_k_tokens.device();
const auto kernel_device = top_k.device();
const auto device = LaunchKernel::resolve_device(kernel_device);
const void* const host_cache_k_ptr = runtime::get_device_accessible_ptr(host_cache_k);
const void* const host_cache_v_ptr =
@@ -741,7 +776,7 @@ void load_cache_to_device_buffer(
// Generic lambda: int32/int64 kernel variants are compiled for both
// seq_lens and req_pool_indices; the correct combo is selected at runtime.
auto launch = [&](auto kernel_fn, const auto* seq_lens_ptr, const auto* req_pool_indices_ptr) {
constexpr size_t smem_bytes = SmemLayout<NUM_TOP_K, HOT_BUFFER_SIZE>::BYTES;
constexpr size_t smem_bytes = SmemLayout<NUM_TOP_K_TOKENS, HOT_BUFFER_SIZE>::BYTES;
#ifndef USE_ROCM
if constexpr (smem_bytes > 48u * 1024u) {
cudaFuncSetAttribute(kernel_fn, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_bytes);
@@ -749,7 +784,7 @@ void load_cache_to_device_buffer(
#endif
LaunchKernel(bs, BLOCK_SIZE, device, smem_bytes)(
kernel_fn,
static_cast<const int32_t*>(top_k_tokens.data_ptr()),
static_cast<const int32_t*>(top_k.data_ptr()),
static_cast<int32_t*>(device_buffer_tokens.data_ptr()),
static_cast<const int64_t*>(host_cache_locs.data_ptr()),
static_cast<const int32_t*>(device_buffer_locs.data_ptr()),
@@ -765,7 +800,7 @@ void load_cache_to_device_buffer(
buffer_stride_0,
host_stride,
lru_slot_stride_0,
top_k_tokens_stride,
top_k_stride,
top_k_device_locs_stride,
page_size,
item_size_bytes,
@@ -788,6 +823,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int64_t,
@@ -802,6 +839,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int64_t,
@@ -816,6 +855,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int32_t,
@@ -830,6 +871,8 @@ void load_cache_to_device_buffer(
HOT_BUFFER_SIZE,
IsMLA,
IsDsv4Layout,
SPARSE_BLOCK_SIZE,
TopKIsBlocks,
RecordMissPlan,
SkipIO,
int32_t,
@@ -23,6 +23,7 @@ from ..common.utils import (
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
"HAS_SINK": lambda args: args["sink_ptr"] is not None,
"BATCH_SIZE_BUCKET": lambda args: triton.next_power_of_2(args["batch_size"]),
"HAS_HISPARSE_SLOTS": lambda args: args["hisparse_slots_ptr"] is not None,
}
)
@triton.autotune(
@@ -43,6 +44,7 @@ def _gqa_share_sparse_decode_kernel(
idx_ptr, # topk index: qh x b x topk
o_ptr, # O partial: c x b x qh x d
lse_ptr, # lse partial: c x b x qh
hisparse_slots_ptr, # pre-resolved device slots: kh x b x (topk * block)
seq_lens,
slot_ids,
# shape
@@ -52,6 +54,8 @@ def _gqa_share_sparse_decode_kernel(
head_dim,
max_topk,
max_kv_len,
hisparse_slots_stride_h,
hisparse_slots_stride_b,
# sm_scale
sm_scale,
# per-tensor KV dequant scales (1.0 when the cache is unit-scaled)
@@ -89,6 +93,7 @@ def _gqa_share_sparse_decode_kernel(
NUM_TOPK_CHUNKS: tl.constexpr,
HAS_SINK: tl.constexpr,
IS_FP8: tl.constexpr,
HAS_HISPARSE_SLOTS: tl.constexpr,
):
# decode program ids: split-K over the topk dimension to give every SM
# something to do at small batch. pid(0) folds (batch, chunk) together so
@@ -161,13 +166,25 @@ def _gqa_share_sparse_decode_kernel(
# only iterate over this chunk's topk slice. the load must respect the
# per-chunk start offset.
cur_idx_ptr = idx_base + chunk_start_topk * stride_ti_t
hisparse_topk_counter = chunk_start_topk
for _ in tl.range(chunk_start_topk, chunk_end_topk):
# load index
c = tl.load(cur_idx_ptr).to(tl.int32) * BLOCK_SIZE_N
cur_idx_ptr = cur_idx_ptr + stride_ti_t
# resolve slots for this block via req_to_token
pos = c + off_n
pos_mask = pos < seq_len
if HAS_HISPARSE_SLOTS:
slots = tl.load(
hisparse_slots_ptr
+ pid_kh * hisparse_slots_stride_h
+ pid_b * hisparse_slots_stride_b
+ hisparse_topk_counter * BLOCK_SIZE_N
+ off_n,
mask=off_n < BLOCK_SIZE_N,
other=0,
).to(tl.int64)
hisparse_topk_counter = hisparse_topk_counter + 1
else:
slots = tl.load(
req_to_token_ptr + sid * stride_r2t_b + pos,
mask=pos_mask,
@@ -321,6 +338,7 @@ def flash_decode_with_gqa_share_sparse(
q_scale: Optional[float] = None,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
hisparse_slots: Optional[torch.Tensor] = None,
) -> torch.Tensor:
triton.set_allocator(robust_allocator)
is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="decode")
@@ -384,6 +402,7 @@ def flash_decode_with_gqa_share_sparse(
topk_idx,
o_partial,
lse_partial,
hisparse_slots,
seq_lens,
slot_ids,
max_slots,
@@ -392,6 +411,8 @@ def flash_decode_with_gqa_share_sparse(
head_dim,
max_topk,
max_kv_len,
hisparse_slots.stride(0) if hisparse_slots is not None else 0,
hisparse_slots.stride(1) if hisparse_slots is not None else 0,
sm_scale,
k_scale,
v_scale,
@@ -28,6 +28,7 @@ from ..common.utils import (
"BLOCK_SIZE_T": lambda args: triton.next_power_of_2(args["max_topk"]),
"BLOCK_SIZE_QH": lambda args: args["BLOCK_SIZE_Q"] * args["BLOCK_SIZE_H"],
"HAS_SINK": lambda args: args["sink_ptr"] is not None,
"HAS_LOC_MAPPING": lambda args: args["loc_mapping_ptr"] is not None,
}
)
@triton.autotune(
@@ -55,6 +56,7 @@ def _gqa_share_sparse_fwd_kernel(
t_ptr, # topk_idx: kh x n x k
o_ptr, # O: n x h x d
req_to_token_ptr, # req_to_token: max_reqs x max_kv_len
loc_mapping_ptr, # logical slot to HiSparse device slot
# seqlens
cu_seqlens_q,
cu_seqblocks_q,
@@ -106,6 +108,7 @@ def _gqa_share_sparse_fwd_kernel(
HAS_SINK: tl.constexpr,
USE_TMA: tl.constexpr,
IS_FP8: tl.constexpr,
HAS_LOC_MAPPING: tl.constexpr,
):
sm_scale_log2e = sm_scale * 1.4426950409
# get batch id and head id
@@ -199,6 +202,12 @@ def _gqa_share_sparse_fwd_kernel(
mask=pos_mask,
other=0,
).to(tl.int64)
if HAS_LOC_MAPPING:
slots = tl.load(
loc_mapping_ptr + slots,
mask=pos_mask,
other=0,
).to(tl.int64)
slots = (slots + max_slots) % max_slots # safety against negative
# k shape: [BLOCK_SIZE_KD, BLOCK_SIZE_K] (transposed for tl.dot)
k = tl.load(
@@ -289,6 +298,7 @@ def flash_prefill_with_gqa_share_sparse(
q_scale: Optional[float] = None,
k_scale: Optional[float] = None,
v_scale: Optional[float] = None,
loc_mapping: Optional[torch.Tensor] = None,
) -> torch.Tensor:
triton.set_allocator(robust_allocator)
is_fp8 = check_sparse_kv_fp8(q, k_cache, v_cache, label="prefill")
@@ -340,6 +350,7 @@ def flash_prefill_with_gqa_share_sparse(
topk_idx,
o,
req_to_token,
loc_mapping,
cu_seqlens,
cu_seqblocks_q,
seq_lens,
+84 -1
View File
@@ -174,6 +174,8 @@ def _jit_sparse_module(
hot_buffer_size: int,
is_mla: bool = False,
is_dsv4_layout: bool = False,
top_k_block_size: int = 1,
top_k_is_blocks: bool = False,
record_miss_plan: bool = False,
skip_io: bool = False,
) -> Module:
@@ -185,6 +187,8 @@ def _jit_sparse_module(
hot_buffer_size,
is_mla,
is_dsv4_layout,
top_k_block_size,
top_k_is_blocks,
record_miss_plan,
skip_io,
)
@@ -195,6 +199,8 @@ def _jit_sparse_module(
hot_buffer_size,
is_mla,
is_dsv4_layout,
top_k_block_size,
top_k_is_blocks,
record_miss_plan,
skip_io,
)
@@ -308,7 +314,7 @@ def _load_cache_to_device_buffer_mla(
skip_io=skip_io,
)
empty = torch.empty(0)
empty = torch.empty(0, device=top_k_tokens.device)
if num_real_reqs is None:
num_real_reqs = torch.tensor(
@@ -399,6 +405,83 @@ def load_cache_to_device_buffer_mla(
)
def load_blocks_to_device_buffer_mha(
top_k_blocks: torch.Tensor,
device_buffer_tokens: torch.Tensor,
host_cache_locs: torch.Tensor,
device_buffer_locs: torch.Tensor,
host_cache_k: torch.Tensor,
host_cache_v: torch.Tensor,
device_buffer_k: torch.Tensor,
device_buffer_v: torch.Tensor,
top_k_device_locs: torch.Tensor,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
lru_slots: torch.Tensor,
item_size_bytes: int,
hot_buffer_size: int,
sparse_block_size: int,
page_size: int = 1,
block_size: int = 256,
num_real_reqs: torch.Tensor | None = None,
skip_io: bool = False,
) -> None:
"""Swap block-selected MHA K/V into the HiSparse device pool."""
num_top_k_blocks = top_k_blocks.size(1)
num_top_k_tokens = num_top_k_blocks * sparse_block_size
assert hot_buffer_size >= num_top_k_tokens, (
f"hot_buffer_size ({hot_buffer_size}) must be >= selected tokens "
f"({num_top_k_tokens})"
)
assert top_k_device_locs.size(1) >= num_top_k_tokens
k_stride = host_cache_k.stride(0) * host_cache_k.element_size()
v_stride = host_cache_v.stride(0) * host_cache_v.element_size()
assert k_stride == v_stride == item_size_bytes, (
"K/V token strides must equal item_size_bytes: "
f"k_stride={k_stride}, v_stride={v_stride}, "
f"item_size_bytes={item_size_bytes}"
)
module = _jit_sparse_module(
item_size_bytes,
block_size,
num_top_k_blocks,
hot_buffer_size,
is_mla=False,
is_dsv4_layout=False,
top_k_block_size=sparse_block_size,
top_k_is_blocks=True,
record_miss_plan=False,
skip_io=skip_io,
)
empty = torch.empty(0, device=top_k_blocks.device)
if num_real_reqs is None:
num_real_reqs = torch.tensor(
[top_k_blocks.size(0)], dtype=torch.int32, device=top_k_blocks.device
)
module.load_cache_to_device_buffer(
top_k_blocks,
device_buffer_tokens,
host_cache_locs,
device_buffer_locs,
host_cache_k,
host_cache_v,
device_buffer_k,
device_buffer_v,
top_k_device_locs,
req_pool_indices,
seq_lens,
lru_slots,
num_real_reqs,
page_size,
item_size_bytes,
empty,
empty,
empty,
)
def copy_cache_planned_mla(
*,
miss_src: torch.Tensor,
@@ -86,14 +86,16 @@ def validate_hisparse(server_args: ServerArgs) -> None:
from sglang.srt.configs.model_config import (
is_deepseek_dsa,
is_deepseek_v4,
is_minimax_sparse,
)
hf_config = model_config_of(server_args).hf_config
is_v4_hisparse = is_deepseek_v4(hf_config)
is_m3_hisparse = is_minimax_sparse(hf_config)
is_hip = get_platform().is_hip
assert is_deepseek_dsa(hf_config) or is_v4_hisparse, (
assert is_deepseek_dsa(hf_config) or is_v4_hisparse or is_m3_hisparse, (
"--enable-hisparse is only supported for DSA (DeepSeek Sparse Attention) "
"models (e.g., DeepSeek V3.2, GLM-5) and DeepSeek V4 now. "
"models (e.g., DeepSeek V3.2, GLM-5), DeepSeek V4, and MiniMax M3 now. "
)
assert cfg.disable_radix_cache, (
@@ -121,6 +123,10 @@ def validate_hisparse(server_args: ServerArgs) -> None:
)
return
# MiniMax M3 uses its own Triton sparse kernels.
if is_m3_hisparse:
return
if resolved_view(server_args).kv_cache_dtype not in (
"bfloat16",
"auto",
@@ -24,6 +24,7 @@ class StateType(str, enum.Enum):
# only the live subrange of that row for the current open pool.
DSA_TAIL = "dsa_tail"
MINIMAX_INDEX_K = "minimax_index_k"
MINIMAX_DENSE_KV = "minimax_dense_kv"
# DeepSeek-V4 unified_kv SWA ring: addressed per-row by ring slot
# (req_pool_idx * ring_stride + pos % ring_stride), needs its own component.
SWA_RING = "swa_ring"
+4 -2
View File
@@ -62,6 +62,7 @@ from sglang.srt.disaggregation.utils import (
build_staging_slot_metadata,
get_dsa_tail_state_indices,
get_kv_class,
get_kv_transfer_buf_infos,
get_qsa_pending_state_indices,
is_mla_backend,
is_unadmitted_reject,
@@ -575,8 +576,8 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
if self.scheduler.enable_hisparse
else self.token_to_kv_pool
)
kv_data_ptrs, kv_data_lens, kv_item_lens = (
transfer_kv_pool.get_contiguous_buf_infos()
kv_data_ptrs, kv_data_lens, kv_item_lens = get_kv_transfer_buf_infos(
transfer_kv_pool
)
kv_data_mem_kinds = (
["DRAM"] * len(kv_data_ptrs)
@@ -1579,6 +1580,7 @@ class DecodePreallocQueue(DecodeHiCachePreallocMixin):
StateType.DSA: _full_kv_pages_payload,
StateType.DSA_TAIL: _dsa_tail_payload,
StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
StateType.MINIMAX_DENSE_KV: _full_kv_pages_payload,
StateType.SWA_RING: _swa_ring_payload,
StateType.DSV4_REQUEST_STATE: _request_state_payload,
StateType.BLOCK_SCALE: _full_kv_pages_payload,
@@ -1398,7 +1398,9 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
if src_layer_ids or dst_layer_ids:
# Draft buffers break the flat [K block, V block] layout, so pair by
# layer ID instead of the half-split used by get_mha_kv_ptrs_with_pp.
if any(l != src_kv_item_len for l in self.kv_args.kv_item_lens):
if any(
item_len != src_kv_item_len for item_len in self.kv_args.kv_item_lens
):
logger.error(
f"[{mooncake_session_id}] head-sliced transfer assumes one item "
f"length for every KV entry, got {set(self.kv_args.kv_item_lens)}"
@@ -1852,12 +1854,11 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
)
or rc
)
elif st == StateType.MINIMAX_INDEX_K:
# Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
# lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
elif st in (StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV):
# Compacted layer lists require equal TP and PP=1 on both peers.
if self.pp_size is not None and self.pp_size > 1:
raise RuntimeError(
"PD disagg: PP>1 not supported for MiniMax sparse index yet."
"PD disagg: PP>1 not supported for MiniMax state yet."
)
if (
target_rank_registration_info is not None
@@ -1866,11 +1867,17 @@ class MooncakeKVManager(StagingManagerMixin, CommonKVManager):
):
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax "
"sparse index yet."
"state yet."
)
src_indices = list(indices)
dst_indices_local = list(dst_indices)
if len(src_indices) > len(dst_indices_local):
if st == StateType.MINIMAX_DENSE_KV:
if len(src_indices) != len(dst_indices_local):
raise RuntimeError(
f"{st.value} state index length mismatch: "
f"prefill={len(src_indices)}, dst={len(dst_indices_local)}"
)
elif len(src_indices) > len(dst_indices_local):
src_indices = src_indices[: len(dst_indices_local)]
elif len(src_indices) < len(dst_indices_local):
dst_indices_local = dst_indices_local[: len(src_indices)]
@@ -1282,6 +1282,7 @@ class MoriKVManager(CommonKVManager):
"swa_ring",
"c128_state",
"minimax_index_k",
"minimax_dense_kv",
):
statuses.extend(
self._send_swa_dsa_state(
@@ -1409,7 +1410,12 @@ class MoriKVManager(CommonKVManager):
f"PD state transfer does not support TP-mismatched non-MLA SWA models "
f"(prefill_tp_size={self.attn_tp_size}, decode_tp_size={peer_info.decode_tp_size})"
)
if state_type in ("qsa_pending", "qsa_compressed", "minimax_index_k"):
if state_type in (
"qsa_pending",
"qsa_compressed",
"minimax_index_k",
"minimax_dense_kv",
):
if self.pp_size is not None and self.pp_size > 1:
# MORI registration does not exchange state_layer_ids. Compact
# sparse-state lists therefore cannot be paired safely across
@@ -1445,6 +1451,7 @@ class MoriKVManager(CommonKVManager):
"qsa_compressed",
"swa_ring",
"c128_state",
"minimax_dense_kv",
):
raise RuntimeError(
f"{state_type.upper()} state index length mismatch: "
@@ -2696,17 +2696,16 @@ class NixlKVManager(StagingManagerMixin, CommonKVManager):
dst_layer_ids=dst_lids,
dst_item_lens=dst_lens,
)
elif st == StateType.MINIMAX_INDEX_K:
# Equal-TP / PP=1 only. Sub-pools are compacted sparse-layer
# lists, so PP>1 mis-slices and heterogeneous TP is unsupported.
elif st in (StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV):
# Compacted layer lists require equal TP and PP=1 on both peers.
if self.pp_size is not None and self.pp_size > 1:
raise RuntimeError(
"PD disagg: PP>1 not supported for MiniMax sparse index yet."
"PD disagg: PP>1 not supported for MiniMax state yet."
)
if self.attn_tp_size != decode_tp_size:
raise RuntimeError(
"PD disagg: heterogeneous TP not supported for MiniMax "
"sparse index yet."
"state yet."
)
if len(src_indices) != len(dst_indices):
raise RuntimeError(
+4 -2
View File
@@ -53,6 +53,7 @@ from sglang.srt.disaggregation.utils import (
build_staging_slot_metadata,
get_dsa_tail_state_indices,
get_kv_class,
get_kv_transfer_buf_infos,
get_qsa_pending_state_indices,
is_aborted,
is_mla_backend,
@@ -256,8 +257,8 @@ class PrefillBootstrapQueue:
hf_text_config=self.scheduler.model_config.hf_text_config,
)
)
kv_data_ptrs, kv_data_lens, kv_item_lens = (
self.token_to_kv_pool.get_contiguous_buf_infos()
kv_data_ptrs, kv_data_lens, kv_item_lens = get_kv_transfer_buf_infos(
self.token_to_kv_pool
)
kv_args.prefill_end_layer = (
kv_args.prefill_start_layer + len(kv_data_ptrs)
@@ -1424,6 +1425,7 @@ class SchedulerDisaggregationPrefillMixin:
StateType.DSA: _full_kv_pages_payload,
StateType.DSA_TAIL: _dsa_tail_payload,
StateType.MINIMAX_INDEX_K: _full_kv_pages_payload,
StateType.MINIMAX_DENSE_KV: _full_kv_pages_payload,
StateType.SWA_RING: _swa_ring_payload,
StateType.DSV4_REQUEST_STATE: _request_state_payload,
StateType.BLOCK_SCALE: _full_kv_pages_payload,
+13
View File
@@ -1310,6 +1310,14 @@ def build_dsa_tail_transfer_blocks(
return transfer_blocks
def get_kv_transfer_buf_infos(pool):
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
if isinstance(pool, MiniMaxSparseKVPool):
return pool.get_sparse_kv_buf_infos()
return pool.get_contiguous_buf_infos()
def setup_state_kv_args(
kv_args: KVArgs,
token_to_kv_pool,
@@ -1375,6 +1383,11 @@ def setup_state_kv_args(
if token_to_kv_pool.index_k_pool is not None:
dp, dl, il = token_to_kv_pool.get_index_k_state_buf_infos()
append_state_component(kv_args, StateType.MINIMAX_INDEX_K, dp, dl, il)
append_state_component(
kv_args,
StateType.MINIMAX_DENSE_KV,
*token_to_kv_pool.get_dense_kv_state_buf_infos(),
)
elif hasattr(token_to_kv_pool, "get_state_buf_infos"):
data_ptrs, data_lens, item_lens = token_to_kv_pool.get_state_buf_infos()
@@ -116,6 +116,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
assert isinstance(runner.token_to_kv_pool, MiniMaxSparseKVPool)
self.is_npu = is_npu()
self.kv_pool = runner.token_to_kv_pool
self.hisparse_coordinator = runner.hisparse_coordinator
self.token_to_kv_pool = runner.token_to_kv_pool # alias for TboAttnBackend
self.req_to_token_pool = runner.req_to_token_pool # pool obj for TboAttnBackend
self.req_to_token = runner.req_to_token_pool.req_to_token
@@ -176,6 +177,18 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
local_tokens + self.block_size_k - 1
) // self.block_size_k + 1
self.topk_blocks = sparse_cfg["sparse_topk_blocks"]
if self.hisparse_coordinator is not None:
selected_tokens = self.topk_blocks * self.block_size_k
assert selected_tokens <= self.hisparse_coordinator.device_buffer_size, (
f"MiniMax M3 selects {selected_tokens} sparse-attention tokens, "
"but the HiSparse device buffer holds only "
f"{self.hisparse_coordinator.device_buffer_size}."
)
self._loc_mapping = (
self.kv_pool.main_pool.full_to_hisparse_device_index_mapping
)
else:
self._loc_mapping = None
# MSA (fmha_sm100) is SM100-only; fall back to the Triton sparse path when
# the kernel is unavailable or its constraints don't hold.
@@ -209,6 +222,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
)
self.use_msa = (
not envs.SGLANG_DISABLE_MSA.get()
and self.hisparse_coordinator is None
and msa_available()
and self.block_size_k == 128
and self.kv_pool.page_size == self.block_size_k
@@ -245,6 +259,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
self.page_size = self.kv_pool.page_size
self.use_dense_sparse_decode = (
(not self.is_npu)
and self.hisparse_coordinator is None
and envs.SGLANG_OPT_USE_MINIMAX_DENSE_SPARSE_DECODE.get()
and self.block_size_k % self.page_size == 0
# _dense_sparse_main_decode calls trtllm decode with a bf16 q and
@@ -326,6 +341,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
f"msa_owns_decode={self._msa_owns_decode}, "
f"decode_cuda_graph={_decode_cuda_graph}, "
f"fp8_attn_gemm={self.fp8_attn_gemm}, "
f"hisparse={'enabled' if self._loc_mapping is not None else 'disabled'}, "
f"npu_native_attn={'on' if (self._native_sparse_ok and _native_attn_enabled()) else 'off'}, "
f"disable_value_layers={sorted(self.disable_value_layer_ids)})"
)
@@ -336,6 +352,22 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
"take minutes; compiles serialize across TP ranks)."
)
def _hisparse_swap_in_blocks(
self,
forward_batch: ForwardBatch,
topk_idx: torch.Tensor,
layer_id: int,
) -> torch.Tensor:
assert topk_idx.size(0) == 1
top_k_device_locs = self.hisparse_coordinator.swap_in_selected_blocks(
req_pool_indices=forward_batch.req_pool_indices,
seq_lens=forward_batch.seq_lens,
top_k_blocks=topk_idx[0],
layer_id=layer_id,
sparse_block_size=self.block_size_k,
)
return top_k_device_locs.unsqueeze(0)
@staticmethod
def _choose_decode_score_max_chunks(batch_size: int) -> int:
"""Score chunk count per graph bucket.
@@ -1549,6 +1581,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
idx_v_scale=layer.idx_v_scale_float,
cached_topk_idx=cached_topk_idx,
return_topk_idx=want_topk,
loc_mapping=self._loc_mapping,
)
if want_topk:
idx_o, o, reduced_topk_idx = result
@@ -1702,6 +1735,16 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
else:
_cached_topk = _topk_buf
hisparse_swap_in_fn = None
if self.hisparse_coordinator is not None:
def hisparse_swap_in_fn(topk_idx):
return self._hisparse_swap_in_blocks(
forward_batch=forward_batch,
topk_idx=topk_idx,
layer_id=layer.layer_id,
)
idx_o, o = minimax_sparse_decode(
q,
None,
@@ -1735,6 +1778,7 @@ class MiniMaxSparseAttnBackend(AttentionBackend):
idx_v_scale=layer.idx_v_scale_float,
cached_topk_idx=_cached_topk,
topk_out=_topk_buf if _want_topk else None,
hisparse_swap_in_fn=hisparse_swap_in_fn,
)
return (
None if idx_o is None else idx_o.reshape(q.shape[0], -1).contiguous(),
@@ -75,6 +75,7 @@ def minimax_sparse_prefill(
idx_v_scale: Optional[float] = None,
cached_topk_idx: Optional[torch.Tensor] = None,
return_topk_idx: bool = False,
loc_mapping: Optional[torch.Tensor] = None,
):
"""Run MiniMax-M3 sparse prefill.
@@ -146,7 +147,7 @@ def minimax_sparse_prefill(
# Step 3: Sparse attention using topk index (main head). The MSA path only
# replaces this step; the indexer above is unchanged. MSA has no attn-sink
# input, so keep the Triton path when sink is present.
if use_msa and sink is None:
if use_msa and sink is None and loc_mapping is None:
from .msa import MSAUnavailableError, msa_sparse_prefill_main
try:
@@ -188,6 +189,7 @@ def minimax_sparse_prefill(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
loc_mapping=loc_mapping,
)
else:
o = flash_prefill_with_gqa_share_sparse(
@@ -210,6 +212,7 @@ def minimax_sparse_prefill(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
loc_mapping=loc_mapping,
)
if return_topk_idx:
return idx_o, o, reduced_topk_idx
@@ -255,6 +258,7 @@ def minimax_sparse_decode(
idx_v_scale: Optional[float] = None,
cached_topk_idx: Optional[torch.Tensor] = None,
topk_out: Optional[torch.Tensor] = None,
hisparse_swap_in_fn: Optional[Callable] = None,
) -> Tuple[torch.Tensor, torch.Tensor]:
# Index top-k sharing for DECODE. A group's source layer passes ``topk_out``
# (a persistent buffer) and publishes its reduced top-k there; the group's
@@ -319,9 +323,12 @@ def minimax_sparse_decode(
f"reduced top-k shape {tuple(topk_idx.shape)}"
)
topk_out.copy_(topk_idx)
hisparse_slots = (
hisparse_swap_in_fn(topk_idx) if hisparse_swap_in_fn is not None else None
)
# Step 3: Sparse attention using topk index (main head). The MSA path
# only replaces this step; keep the Triton path when sink is present.
if use_msa and sink is None:
if use_msa and sink is None and hisparse_slots is None:
from .msa import MSAUnavailableError, msa_sparse_decode_main
try:
@@ -357,6 +364,7 @@ def minimax_sparse_decode(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
hisparse_slots=hisparse_slots,
)
else:
o = flash_decode_with_gqa_share_sparse(
@@ -373,5 +381,6 @@ def minimax_sparse_decode(
q_scale=q_scale,
k_scale=k_scale,
v_scale=v_scale,
hisparse_slots=hisparse_slots,
)
return idx_o, o
@@ -19,9 +19,16 @@ if is_xpu():
"copy_cache_planned_mla has no AOT sgl_kernel implementation."
)
def load_blocks_to_device_buffer_mha(*args, **kwargs):
raise RuntimeError(
"MiniMax M3 HiSparse block swap-in is unsupported on XPU: "
"load_blocks_to_device_buffer_mha has no AOT sgl_kernel implementation."
)
else:
from sglang.kernels.ops.kvcache.hisparse import (
copy_cache_planned_mla,
load_blocks_to_device_buffer_mha,
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
)
@@ -36,8 +43,9 @@ from sglang.srt.mem_cache.allocator.hisparse import (
from sglang.srt.mem_cache.hisparse_memory_pool import (
HiSparseDSATokenToKVPool,
)
from sglang.srt.mem_cache.memory_pool import ReqToTokenPool
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool, ReqToTokenPool
from sglang.srt.mem_cache.memory_pool_host import DeepSeekV4PagedHostPool
from sglang.srt.mem_cache.pool_host.mha import HiSparseMHATokenToKVPoolHost
from sglang.srt.mem_cache.pool_host.mla import MLATokenToKVPoolHost
device_module = get_device_module()
@@ -157,9 +165,11 @@ class HiSparseCoordinator:
)
self.compress_ratio = self.token_to_kv_pool_allocator.compress_ratio
kvcache = self.token_to_kv_pool_allocator.get_kvcache()
self.is_dsv4_hisparse = isinstance(
self.token_to_kv_pool_allocator, DeepSeekV4HiSparseTokenToKVPoolAllocator
)
self.is_m3_hisparse = isinstance(kvcache, MiniMaxSparseKVPool)
if self.is_dsv4_hisparse:
self.mem_pool_device = self.token_to_kv_pool_allocator.hisparse_kvcache
page_size = self.mem_pool_device.page_size
@@ -184,9 +194,21 @@ class HiSparseCoordinator:
assert isinstance(
self.token_to_kv_pool_allocator, HiSparseTokenToKVPoolAllocator
)
self.mem_pool_device: HiSparseDSATokenToKVPool = (
self.token_to_kv_pool_allocator.get_kvcache()
if self.is_m3_hisparse:
self.mem_pool_device = kvcache.main_pool
assert self.mem_pool_device.head_num == 1, (
"MiniMax M3 HiSparse requires one KV head per TP rank, "
f"got {self.mem_pool_device.head_num}. Increase the "
"tensor-parallel size."
)
self.mem_pool_host = HiSparseMHATokenToKVPoolHost(
device_pool=self.mem_pool_device,
host_to_device_ratio=host_to_device_ratio,
page_size=self.mem_pool_device.page_size,
)
self.item_size_bytes = self.mem_pool_device.bytes_per_token_k
else:
self.mem_pool_device: HiSparseDSATokenToKVPool = kvcache
self.mem_pool_host = MLATokenToKVPoolHost(
device_pool=self.mem_pool_device,
host_to_device_ratio=host_to_device_ratio,
@@ -263,9 +285,17 @@ class HiSparseCoordinator:
self.device_buffer_size, dtype=torch.int32, device=device
)
# Pre-allocated output buffer for swap_in_selected_pages (CUDA-graph safe)
# Pre-allocated output buffer for swap-in (CUDA-graph safe). MiniMax
# selects blocks, so its flattened token-slot output can occupy any
# prefix up to the full device working-set size.
swap_output_width = (
self.device_buffer_size if self.is_m3_hisparse else self.top_k
)
self.top_k_device_locs_buffer = torch.full(
(max_num_req_slots, self.top_k), -1, dtype=torch.int32, device=device
(max_num_req_slots, swap_output_width),
-1,
dtype=torch.int32,
device=device,
)
self.raw_indices_buffer = torch.full(
(max_num_req_slots, self.top_k), -1, dtype=torch.int32, device=device
@@ -457,7 +487,12 @@ class HiSparseCoordinator:
host_indices = self.req_to_host_pool[req.kv.req_pool_idx, :n]
device_locs = self.req_to_device_buffer[req.kv.req_pool_idx, :n]
for layer_id in range(self.mem_pool_device.layer_num):
layer_ids = (
range(self.mem_pool_device.start_layer, self.mem_pool_device.end_layer)
if self.is_m3_hisparse
else range(self.mem_pool_device.layer_num)
)
for layer_id in layer_ids:
self.mem_pool_host.load_to_device_per_layer(
self.mem_pool_device,
host_indices,
@@ -642,13 +677,9 @@ class HiSparseCoordinator:
compressed_locs = self.token_to_kv_pool_allocator.get_last_loc_compressed(
out_cache_loc
)
# ROCm: the decode remap creates a temporary hisparse device slot per
# new token (via the page_size==1 allocator path). Free the stale
# slot before pointing the mapping at the reserved device-buffer slot,
# otherwise the temporary slots leak and corrupt later swap-in lookups.
# CUDA keeps the original behavior: the swap-in kernel consumes only
# top_k_device_locs, so stale mapping entries are harmless there.
if _is_hip:
# Page-size-one allocation creates a temporary slot before remapping
# the new token into the request's reserved device-buffer slot.
if _is_hip or self.mem_pool_device.page_size == 1:
previous_locs = self.mem_pool_device._translate_loc_to_hisparse_device(
compressed_locs
)
@@ -976,8 +1007,7 @@ class HiSparseCoordinator:
miss plan into self._miss_{src,dst,count} for the skip layers to replay.
"""
num_reqs = req_pool_indices.size(0)
top_k_indices = self.top_k_device_locs_buffer[:num_reqs]
top_k_indices = self.top_k_device_locs_buffer[:num_reqs, : self.top_k]
swap_in_fn = (
load_cache_to_device_buffer_dsv4_mla
if self.is_dsv4_hisparse
@@ -1015,6 +1045,46 @@ class HiSparseCoordinator:
)
return top_k_indices
def swap_in_selected_blocks(
self,
req_pool_indices: torch.Tensor,
seq_lens: torch.Tensor,
top_k_blocks: torch.Tensor,
layer_id: int,
sparse_block_size: int,
) -> torch.Tensor:
assert self.is_m3_hisparse
num_reqs = req_pool_indices.size(0)
num_selected_tokens = top_k_blocks.size(1) * sparse_block_size
assert num_selected_tokens <= self.device_buffer_size, (
f"MiniMax M3 selected {num_selected_tokens} tokens, but the "
f"HiSparse device buffer holds only {self.device_buffer_size}."
)
top_k_indices = self.top_k_device_locs_buffer[:num_reqs, :num_selected_tokens]
host_layer = layer_id - self.mem_pool_device.start_layer
load_blocks_to_device_buffer_mha(
top_k_blocks=top_k_blocks,
device_buffer_tokens=self.req_device_buffer_tokens[host_layer],
host_cache_locs=self.req_to_host_pool,
device_buffer_locs=self.req_device_buffer_token_locs[host_layer],
host_cache_k=self.mem_pool_host.k_buffer[host_layer],
host_cache_v=self.mem_pool_host.v_buffer[host_layer],
device_buffer_k=self.mem_pool_device.get_key_buffer(layer_id),
device_buffer_v=self.mem_pool_device.get_value_buffer(layer_id),
top_k_device_locs=top_k_indices,
req_pool_indices=req_pool_indices,
seq_lens=seq_lens,
lru_slots=self.lru_slots[host_layer],
item_size_bytes=self.item_size_bytes,
hot_buffer_size=self.device_buffer_size,
sparse_block_size=sparse_block_size,
page_size=1,
block_size=self.swap_in_block_size,
num_real_reqs=self.num_real_reqs,
skip_io=self.skip_io,
)
return top_k_indices
def _run_copy_only_kernel(self, num_reqs: int, skip_layer: int) -> None:
"""Replay the anchor's recorded miss plan into a skip layer's buffers
(IO-only; the anchor's slot table stays valid -- lockstep layout)."""
@@ -1045,7 +1115,10 @@ class HiSparseCoordinator:
"""
if not self.enable_prefetch:
return self._run_swap_in_kernel(
req_pool_indices, compressed_seq_lens, top_k_result, layer_id
req_pool_indices,
compressed_seq_lens,
top_k_result,
layer_id,
)
num_reqs = req_pool_indices.size(0)
@@ -1054,7 +1127,7 @@ class HiSparseCoordinator:
# applies (shared index + lockstep buffers).
slot = self._prefetch_slot[layer_id]
self._prefetch_events[slot].wait(device_module.current_stream())
return self.top_k_device_locs_buffer[:num_reqs]
return self.top_k_device_locs_buffer[:num_reqs, : self.top_k]
# Anchor: swap in synchronously (recording the plan), then prefetch the
# skip layers' copies on the side stream.
@@ -1,4 +1,7 @@
from __future__ import annotations
import weakref
from typing import TYPE_CHECKING
import torch
@@ -11,6 +14,9 @@ from sglang.srt.mem_cache.deepseek_v4_memory_pool import (
from sglang.srt.mem_cache.hisparse_memory_pool import HiSparseDSATokenToKVPool
from sglang.srt.utils.common import get_num_new_pages
if TYPE_CHECKING:
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
def __init__(
@@ -19,7 +25,7 @@ class HiSparseTokenToKVPoolAllocator(BaseTokenToKVPoolAllocator):
page_size: int,
dtype: torch.dtype,
device: torch.device,
kvcache: HiSparseDSATokenToKVPool,
kvcache: HiSparseDSATokenToKVPool | MiniMaxSparseKVPool,
need_sort: bool,
host_to_device_ratio: int = 2,
):
@@ -9,7 +9,7 @@ from sglang.kernels.ops.kvcache.hisparse_slot_mapping import (
translate_padded_hisparse_locations,
)
from sglang.srt.layers.radix_attention import RadixAttention
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool
from sglang.srt.mem_cache.memory_pool import DSATokenToKVPool, MHATokenToKVPool
from sglang.srt.utils import is_cuda, is_hip, is_xpu
logger = logging.getLogger(__name__)
@@ -147,3 +147,111 @@ class HiSparseDSATokenToKVPool(DSATokenToKVPool):
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
raise NotImplementedError("HiSparseDevicePool does not support load_cpu_copy")
class HiSparseMHAMainPool(MHATokenToKVPool):
"""MHA KV pool with HiSparse logical-to-device mapping.
Used by MiniMax M3 HiSparse. The index pools (index_kv_pool, index_k_pool)
stay fully resident on the device and do not use this mapping.
"""
def __init__(
self,
size: int,
page_size: int,
dtype: torch.dtype,
head_num: int,
head_dim: int,
layer_num: int,
device: str,
enable_memory_saver: bool,
start_layer: Optional[int] = None,
end_layer: Optional[int] = None,
):
super().__init__(
size=size,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
layer_num=layer_num,
device=device,
enable_memory_saver=enable_memory_saver,
start_layer=start_layer,
end_layer=end_layer,
)
self.full_to_hisparse_device_index_mapping: Optional[torch.Tensor] = None
self.bytes_per_token_k = head_num * head_dim * self.store_dtype.itemsize
self.bytes_per_token_v = head_num * self.v_head_dim * self.store_dtype.itemsize
def register_mapping(
self, full_to_hisparse_device_index_mapping: torch.Tensor
) -> None:
self.full_to_hisparse_device_index_mapping = (
full_to_hisparse_device_index_mapping
)
def translate_loc_to_hisparse_device(self, indices: torch.Tensor) -> torch.Tensor:
assert self.full_to_hisparse_device_index_mapping is not None
return self.full_to_hisparse_device_index_mapping[indices]
def _translate_loc_to_hisparse_device(self, indices: torch.Tensor) -> torch.Tensor:
assert self.full_to_hisparse_device_index_mapping is not None
return self.full_to_hisparse_device_index_mapping[indices]
def translate_loc_from_full_to_hisparse_device(
self, full_indices: torch.Tensor
) -> torch.Tensor:
assert self.full_to_hisparse_device_index_mapping is not None
return self.full_to_hisparse_device_index_mapping[full_indices]
def translate_loc_from_full_to_compressed(
self, full_indices: torch.Tensor
) -> torch.Tensor:
return full_indices
def set_kv_buffer(
self,
layer: RadixAttention,
loc,
cache_k: torch.Tensor,
cache_v: torch.Tensor,
*args,
**kwargs,
):
from sglang.srt.mem_cache.memory_pool import unwrap_write_loc
raw_loc, _, _ = unwrap_write_loc(loc)
translated = self.translate_loc_to_hisparse_device(raw_loc)
super().set_kv_buffer(layer, translated, cache_k, cache_v, *args, **kwargs)
def transfer_values_on_device(
self,
dst_indices: torch.Tensor,
src_indices: torch.Tensor,
) -> None:
transfer_kv_all_layer_mla(
src_layers=self.k_data_ptrs,
dst_layers=self.k_data_ptrs,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=self.bytes_per_token_k,
num_layers=self.layer_num,
)
transfer_kv_all_layer_mla(
src_layers=self.v_data_ptrs,
dst_layers=self.v_data_ptrs,
src_indices=src_indices,
dst_indices=dst_indices,
item_size=self.bytes_per_token_v,
num_layers=self.layer_num,
)
def get_cpu_copy(self, indices, mamba_indices=None, req_pool_index=None):
raise NotImplementedError("HiSparseMHAMainPool does not support get_cpu_copy")
def load_cpu_copy(
self, kv_cache_cpu, indices, mamba_indices=None, req_pool_index=None
):
raise NotImplementedError("HiSparseMHAMainPool does not support load_cpu_copy")
@@ -1836,6 +1836,14 @@ class KVCacheConfigurator:
disable_value_sparse_layer_ids = get_minimax_sparse_disable_value_layer_ids(
sparse_cfg
)
enable_hisparse = get_memory().enable_hisparse
hisparse_kwargs = {}
if enable_hisparse:
from sglang.srt.mem_cache.sparsity import parse_hisparse_config
hisparse_kwargs["host_to_device_ratio"] = (
parse_hisparse_config().host_to_device_ratio
)
token_to_kv_pool = MiniMaxSparseKVPool(
size=max_total_num_tokens,
page_size=self.pool_page_size,
@@ -1861,6 +1869,8 @@ class KVCacheConfigurator:
enable_memory_saver=get_exec().features.enable_memory_saver,
start_layer=self.layer_info.start_layer,
end_layer=self.layer_info.end_layer,
enable_hisparse=enable_hisparse,
**hisparse_kwargs,
)
return token_to_kv_pool
+123 -13
View File
@@ -5448,6 +5448,8 @@ class MiniMaxSparseKVPool(KVCache):
main_pool_cls=MHATokenToKVPool,
index_kv_pool_cls=MHATokenToKVPool,
index_k_pool_cls=MHATokenToKOnlyPool,
enable_hisparse: bool = False,
host_to_device_ratio: int = 2,
):
# Do not call super().__init__() — delegate to sub-pools instead.
self.size = size
@@ -5466,6 +5468,7 @@ class MiniMaxSparseKVPool(KVCache):
]
index_dtype = index_dtype if index_dtype is not None else dtype
index_pool_size = size * host_to_device_ratio if enable_hisparse else size
# Split sparse layers by V policy: kv_sparse (index_kv_pool holds K+V) vs
# k_only_sparse (index_k_pool holds only K; V is never read).
@@ -5489,13 +5492,50 @@ class MiniMaxSparseKVPool(KVCache):
gid: i for i, gid in enumerate(local_k_only_sparse_layer_ids)
}
self._dense_layer_ids = set(local_dense_layer_ids)
main_layer_num = len(local_dense_layer_ids) + len(local_sparse_layer_ids)
if enable_hisparse:
from sglang.srt.mem_cache.hisparse_memory_pool import (
HiSparseMHAMainPool,
)
self.dense_pool = (
main_pool_cls(
size=index_pool_size,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
layer_num=len(local_dense_layer_ids),
device=device,
enable_memory_saver=enable_memory_saver,
start_layer=start_layer,
end_layer=start_layer + len(local_dense_layer_ids),
)
if local_dense_layer_ids
else None
)
self.main_pool = HiSparseMHAMainPool(
size=size,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
layer_num=len(local_sparse_layer_ids),
device=device,
enable_memory_saver=enable_memory_saver,
start_layer=local_sparse_layer_ids[0],
end_layer=end_layer,
)
else:
self.dense_pool = None
self.main_pool = main_pool_cls(
size=size,
page_size=page_size,
dtype=dtype,
head_num=head_num,
head_dim=head_dim,
layer_num=len(local_dense_layer_ids) + len(local_sparse_layer_ids),
layer_num=main_layer_num,
device=device,
enable_memory_saver=enable_memory_saver,
start_layer=start_layer,
@@ -5504,7 +5544,7 @@ class MiniMaxSparseKVPool(KVCache):
self.index_kv_pool: Optional[MHATokenToKVPool] = (
index_kv_pool_cls(
size=size,
size=index_pool_size,
page_size=page_size,
dtype=index_dtype,
head_num=1,
@@ -5519,7 +5559,7 @@ class MiniMaxSparseKVPool(KVCache):
self.index_k_pool: Optional[MHATokenToKOnlyPool] = (
index_k_pool_cls(
size=size,
size=index_pool_size,
page_size=page_size,
dtype=index_dtype,
head_num=1,
@@ -5533,19 +5573,58 @@ class MiniMaxSparseKVPool(KVCache):
)
self.mem_usage = self.main_pool.mem_usage
if self.dense_pool is not None:
self.mem_usage += self.dense_pool.mem_usage
if self.index_kv_pool is not None:
self.mem_usage += self.index_kv_pool.mem_usage
if self.index_k_pool is not None:
self.mem_usage += self.index_k_pool.mem_usage
# HiCacheController reads these from the top-level KV pool wrapper.
self.layer_num = self.main_pool.layer_num
self.start_layer = self.main_pool.start_layer
self.end_layer = self.main_pool.end_layer
self.layer_num = main_layer_num
self.start_layer = start_layer
self.end_layer = end_layer
# PD disaggregation reads these directly (no fallback) off the wrapper.
self.head_num = self.main_pool.head_num
self.head_dim = self.main_pool.head_dim
self.v_head_dim = self.main_pool.v_head_dim
self.store_dtype = self.main_pool.store_dtype
self.layer_transfer_counter = None
self._enable_hisparse = enable_hisparse
def register_mapping(self, mapping: torch.Tensor) -> None:
assert self._enable_hisparse
self.main_pool.register_mapping(mapping)
def _translate_loc_to_hisparse_device(self, indices: torch.Tensor) -> torch.Tensor:
assert self._enable_hisparse
return self.main_pool._translate_loc_to_hisparse_device(indices)
def translate_loc_to_hisparse_device(self, indices: torch.Tensor) -> torch.Tensor:
assert self._enable_hisparse
return self.main_pool.translate_loc_to_hisparse_device(indices)
def translate_loc_from_full_to_hisparse_device(
self, indices: torch.Tensor
) -> torch.Tensor:
assert self._enable_hisparse
return self.main_pool.translate_loc_from_full_to_hisparse_device(indices)
def translate_loc_from_full_to_compressed(
self, indices: torch.Tensor
) -> torch.Tensor:
assert self._enable_hisparse
return self.main_pool.translate_loc_from_full_to_compressed(indices)
@property
def bytes_per_token_k(self) -> int:
assert self._enable_hisparse
return self.main_pool.bytes_per_token_k
@property
def full_to_hisparse_device_index_mapping(self):
assert self._enable_hisparse
return self.main_pool.full_to_hisparse_device_index_mapping
def register_layer_transfer_counter(
self, layer_transfer_counter: LayerDoneCounter
@@ -5561,17 +5640,22 @@ class MiniMaxSparseKVPool(KVCache):
if self.layer_transfer_counter is not None:
self.layer_transfer_counter.wait_until(layer_id - self.start_layer)
def _pool_for(self, layer_id: int) -> MHATokenToKVPool:
if self.dense_pool is not None and layer_id in self._dense_layer_ids:
return self.dense_pool
return self.main_pool
def get_key_buffer(self, layer_id: int) -> torch.Tensor:
self._wait_for_layer(layer_id)
return self.main_pool.get_key_buffer(layer_id)
return self._pool_for(layer_id).get_key_buffer(layer_id)
def get_value_buffer(self, layer_id: int) -> torch.Tensor:
self._wait_for_layer(layer_id)
return self.main_pool.get_value_buffer(layer_id)
return self._pool_for(layer_id).get_value_buffer(layer_id)
def get_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
self._wait_for_layer(layer_id)
return self.main_pool.get_kv_buffer(layer_id)
return self._pool_for(layer_id).get_kv_buffer(layer_id)
def get_index_kv_buffer(self, layer_id: int) -> Tuple[torch.Tensor, torch.Tensor]:
self._wait_for_layer(layer_id)
@@ -5613,7 +5697,7 @@ class MiniMaxSparseKVPool(KVCache):
Scale semantics follow MHATokenToKVPool: None means unit scale;
a non-None scale is applied with an in-place div_ before the fp8 cast.
"""
self.main_pool.set_kv_buffer(
self._pool_for(layer.layer_id).set_kv_buffer(
layer,
loc,
cache_k,
@@ -5711,8 +5795,10 @@ class MiniMaxSparseKVPool(KVCache):
disable_value = cache_idx_v is None
index_pool = self.index_k_pool if disable_value else self.index_kv_pool
if index_pool is not None and self._can_fuse_kv_index_store(
index_pool, cache_k, cache_idx_k
if (
index_pool is not None
and not self._enable_hisparse
and self._can_fuse_kv_index_store(index_pool, cache_k, cache_idx_k)
):
from sglang.kernels.ops.kvcache.minimax_store_kv_index import store_kv_index
@@ -5759,7 +5845,12 @@ class MiniMaxSparseKVPool(KVCache):
)
def get_kv_size_bytes(self):
sub_pools = [self.main_pool, self.index_kv_pool, self.index_k_pool]
sub_pools = [
self.main_pool,
self.dense_pool,
self.index_kv_pool,
self.index_k_pool,
]
sizes = [p.get_kv_size_bytes() for p in sub_pools if p is not None]
return sum(k for k, _ in sizes), sum(v for _, v in sizes)
@@ -5767,6 +5858,25 @@ class MiniMaxSparseKVPool(KVCache):
# Main K/V only; index buffers ride the state-buffer channel.
return self.main_pool.get_contiguous_buf_infos()
def get_sparse_kv_buf_infos(self):
return self._get_layer_kv_buf_infos(
layer_ids=sorted(self.sparse_layer_id_mapping)
)
def get_dense_kv_state_buf_infos(self):
# Dense KV uses logical device slots, independently of sparse host slots.
return self._get_layer_kv_buf_infos(layer_ids=sorted(self._dense_layer_ids))
def _get_layer_kv_buf_infos(self, *, layer_ids):
buffers = [self.get_key_buffer(layer_id) for layer_id in layer_ids] + [
self.get_value_buffer(layer_id) for layer_id in layer_ids
]
return (
[buffer.data_ptr() for buffer in buffers],
[buffer.nbytes for buffer in buffers],
[buffer[0].nbytes * self.page_size for buffer in buffers],
)
def get_index_k_state_buf_infos(self):
# Per-page item_len (MHATokenToKVPool convention); index rows share the
# main-KV `loc`, so the transfer reuses the same page-ids.
@@ -43,6 +43,7 @@ from sglang.srt.mem_cache.pool_host.common import (
get_allocator_from_storage,
make_kernel_ptr_table,
)
from sglang.srt.mem_cache.pool_host.hisparse import HiSparseHostPoolMixin
from sglang.srt.utils import is_cuda, is_hip, is_mps, is_npu, is_xpu
_is_cuda = is_cuda()
@@ -1098,6 +1099,76 @@ class MHATokenToKOnlyPoolHost(HostKVCache):
return ptr_list, element_size_list
class HiSparseMHATokenToKVPoolHost(HiSparseHostPoolMixin, MHATokenToKVPoolHost):
"""Layer-first MHA host pool with page-granular HiSparse allocation."""
def __init__(
self,
device_pool: MHATokenToKVPool,
host_to_device_ratio: float,
page_size: int,
):
super().__init__(
device_pool=device_pool,
host_to_device_ratio=host_to_device_ratio,
host_size=0,
page_size=page_size,
layout="layer_first",
)
def get_contiguous_buf_infos(self):
buffers = self.k_data_refs + self.v_data_refs
return (
[buffer.data_ptr() for buffer in buffers],
[buffer.nbytes for buffer in buffers],
[buffer[0].nbytes * self.page_size for buffer in buffers],
)
def load_to_device_per_layer(
self,
device_pool,
host_indices,
device_indices,
layer_id,
io_backend,
*,
is_draft: bool = False,
):
if io_backend != "kernel" or is_draft:
raise ValueError(
"MiniMax M3 HiSparse host transfers require the kernel backend."
)
host_layer = layer_id - device_pool.start_layer
transfer_kv_per_layer(
src_k=self.k_buffer[host_layer],
dst_k=device_pool.get_key_buffer(layer_id),
src_v=self.v_buffer[host_layer],
dst_v=device_pool.get_value_buffer(layer_id),
src_indices=host_indices,
dst_indices=device_indices,
item_size=self.token_stride_size,
)
def backup_from_device_all_layer(
self, device_pool, host_indices, device_indices, io_backend
):
if io_backend != "kernel":
raise ValueError(
"MiniMax M3 HiSparse host transfers require the kernel backend."
)
for layer_id in range(device_pool.start_layer, device_pool.end_layer):
host_layer = layer_id - device_pool.start_layer
transfer_kv_per_layer(
src_k=device_pool.get_key_buffer(layer_id),
dst_k=self.k_buffer[host_layer],
src_v=device_pool.get_value_buffer(layer_id),
dst_v=self.v_buffer[host_layer],
src_indices=device_indices,
dst_indices=host_indices,
item_size=self.token_stride_size,
)
class AsymmetricMHATokenToKVPoolHost(MHATokenToKVPoolHost):
"""Host KV pool for MHA models whose K and V have different head dims
(``head_dim != v_head_dim``), e.g. MiMo-V2.
@@ -28,6 +28,7 @@ if is_xpu():
)
else:
from sglang.kernels.ops.kvcache.hisparse import (
load_blocks_to_device_buffer_mha,
load_cache_to_device_buffer_dsv4_mla,
load_cache_to_device_buffer_mla,
transfer_cache_dsv4_mla,
@@ -368,6 +369,84 @@ def test_load_cache_to_device_buffer_hits_newest_and_updates_lru() -> None:
)
@pytest.mark.skipif(is_xpu(), reason="MiniMax MHA block swap-in has no XPU kernel.")
def test_load_blocks_to_device_buffer_mha_handles_partial_newest_block() -> None:
"""A partial newest block must not consume slots for its invalid tail."""
sparse_block_size = 4
hot_buffer_size = 8
host_k = _host_cache()
host_v = _host_cache()
host_v.add_(1000)
device_k = torch.full(
(DEVICE_CACHE_SIZE, 1, KV_DIM), -1, dtype=DTYPE, device=DEVICE
)
device_v = torch.full_like(device_k, -1)
device_buffer_locs = torch.arange(
hot_buffer_size + 1, dtype=torch.int32, device=DEVICE
).view(1, -1)
device_buffer_tokens = torch.tensor(
[[0, 1, 2, 3, -1, -1, -1, -1, -1]],
dtype=torch.int32,
device=DEVICE,
)
for slot, token in enumerate([0, 1, 2, 3]):
device_k[device_buffer_locs[0, slot]].copy_(host_k[token], non_blocking=True)
device_v[device_buffer_locs[0, slot]].copy_(host_v[token], non_blocking=True)
device_k[device_buffer_locs[0, hot_buffer_size]].copy_(
host_k[10], non_blocking=True
)
device_v[device_buffer_locs[0, hot_buffer_size]].copy_(
host_v[10], non_blocking=True
)
top_k_blocks = torch.tensor([[0, 2]], dtype=torch.int32, device=DEVICE)
out = torch.full(
(1, top_k_blocks.size(1) * sparse_block_size),
-1,
dtype=torch.int32,
device=DEVICE,
)
lru_slots = torch.arange(hot_buffer_size, dtype=torch.int16, device=DEVICE).view(
1, -1
)
load_blocks_to_device_buffer_mha(
top_k_blocks=top_k_blocks,
device_buffer_tokens=device_buffer_tokens,
host_cache_locs=torch.arange(
HOST_CACHE_SIZE, dtype=torch.int64, device=DEVICE
).view(1, -1),
device_buffer_locs=device_buffer_locs,
host_cache_k=host_k,
host_cache_v=host_v,
device_buffer_k=device_k,
device_buffer_v=device_v,
top_k_device_locs=out,
req_pool_indices=torch.tensor([0], dtype=torch.int64, device=DEVICE),
seq_lens=torch.tensor([11], dtype=torch.int32, device=DEVICE),
lru_slots=lru_slots,
item_size_bytes=ITEM_SIZE_BYTES,
hot_buffer_size=hot_buffer_size,
sparse_block_size=sparse_block_size,
num_real_reqs=torch.tensor([1], dtype=torch.int32, device=DEVICE),
)
get_device_module().synchronize()
assert torch.equal(
out.cpu(), torch.tensor([[0, 1, 2, 3, 4, 5, 8, -1]], dtype=torch.int32)
)
assert torch.equal(device_k[4].cpu(), host_k[8])
assert torch.equal(device_v[4].cpu(), host_v[8])
assert torch.equal(device_k[5].cpu(), host_k[9])
assert torch.equal(device_v[5].cpu(), host_v[9])
assert torch.equal(
device_buffer_tokens.cpu(),
torch.tensor([[0, 1, 2, 3, 8, 9, -1, -1, -1]], dtype=torch.int32),
)
assert torch.equal(
lru_slots.cpu(), torch.tensor([[6, 7, 4, 5, 0, 1, 2, 3]], dtype=torch.int16)
)
def test_load_cache_to_device_buffer_miss_uses_updated_lru_slot() -> None:
state = _long_case()
@@ -53,14 +53,19 @@ def _make_kv_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
class TestMiniMaxSparseDisaggStateKvArgs(unittest.TestCase):
def test_setup_state_kv_args_single_minimax_component(self):
def test_setup_state_kv_args_minimax_components(self):
pool = _make_k_only_pool()
kv_args = KVArgs()
setup_state_kv_args(kv_args, pool)
self.assertEqual(kv_args.state_types, [StateType.MINIMAX_INDEX_K])
self.assertEqual(len(kv_args.state_data_ptrs), 1)
self.assertEqual(
kv_args.state_types,
[StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV],
)
self.assertEqual(len(kv_args.state_data_ptrs), 2)
self.assertEqual(len(kv_args.state_data_ptrs[0]), pool.index_k_pool.layer_num)
self.assertEqual(len(kv_args.state_item_lens[0]), pool.index_k_pool.layer_num)
self.assertEqual(len(kv_args.state_data_ptrs[1]), 6)
self.assertEqual(len(kv_args.state_item_lens[1]), 6)
def test_index_kv_pool_raises(self):
pool = _make_kv_pool()
@@ -1,4 +1,5 @@
import concurrent.futures
import ctypes
import unittest
from threading import Event
from types import SimpleNamespace
@@ -6,6 +7,7 @@ from unittest.mock import MagicMock, call, patch
import numpy as np
from sglang.srt.disaggregation.base.conn import StateType
from sglang.srt.disaggregation.mooncake.conn import MooncakeKVManager
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
@@ -132,6 +134,61 @@ class TestMooncakeTransferBatching(unittest.TestCase):
)
class TestMiniMaxStateTransfer(CustomTestCase):
def test_index_truncates_but_dense_rejects_mismatched_page_lists(self):
"""Legacy index transfers copy the common prefix; incomplete dense KV must fail."""
def copy_bytes(session, sources, destinations, lengths):
for src, dst, length in zip(sources, destinations, lengths, strict=True):
ctypes.memmove(dst, src, length)
return 0
for state in (StateType.MINIMAX_INDEX_K, StateType.MINIMAX_DENSE_KV):
for src_pages, dst_pages in (([1], [0]), ([1, 2], [0]), ([1], [0, 2])):
with self.subTest(state=state, src=src_pages, dst=dst_pages):
src = np.arange(3, dtype=np.int32)
dst = np.full(3, -1, dtype=np.int32)
manager = MooncakeKVManager.__new__(MooncakeKVManager)
manager.kv_args = SimpleNamespace(
state_types=[state],
state_data_ptrs=[[src.ctypes.data]],
state_item_lens=[[src.itemsize]],
state_dim_per_tensor=[[]],
state_layer_ids=[[]],
)
manager.engine = SimpleNamespace(batch_transfer_sync=copy_bytes)
manager.pp_size = manager.attn_tp_size = 1
manager.is_mla_backend = manager.is_hybrid_mla_backend = False
manager.enable_custom_mem_pool = False
manager.max_transfer_batch_indices = 0
peer = SimpleNamespace(
dst_state_data_ptrs=[[dst.ctypes.data]],
dst_state_item_lens=[[dst.itemsize]],
dst_state_dim_per_tensor=[[]],
dst_state_layer_ids=[[]],
dst_attn_tp_size=1,
)
kwargs = dict(
req=SimpleNamespace(
mooncake_session_id="cpu", dst_state_indices=[dst_pages]
),
prefill_state_indices=[src_pages],
executor=None,
target_rank_registration_info=peer,
)
if state == StateType.MINIMAX_DENSE_KV and len(src_pages) != len(
dst_pages
):
with self.assertRaisesRegex(
RuntimeError, "state index length mismatch"
):
manager.maybe_send_extra(**kwargs)
np.testing.assert_array_equal(dst, [-1, -1, -1])
else:
self.assertEqual(manager.maybe_send_extra(**kwargs), 0)
np.testing.assert_array_equal(dst, [1, -1, -1])
class TestDcpDraftHeadTransfer(unittest.TestCase):
def test_transfers_draft_heads_to_logical_destination_rows(self):
for src_tp, dst_tp in ((4, 8), (8, 4), (8, 8), (4, 32), (32, 4)):
@@ -1,7 +1,7 @@
import unittest
from concurrent.futures import ThreadPoolExecutor
from types import SimpleNamespace
from unittest.mock import MagicMock
from unittest.mock import MagicMock, patch
import numpy as np
import torch
@@ -18,6 +18,71 @@ from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
class TestHiSparseDecodeRemap(CustomTestCase):
def test_page_size_one_reclaims_temporary_device_slot(self):
"""Decode remapping must reclaim its temporary slot without freeing the live slot."""
from sglang.srt.managers.hisparse_coordinator import HiSparseCoordinator
from sglang.srt.mem_cache.allocator.hisparse import (
HiSparseTokenToKVPoolAllocator,
)
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
pool = MiniMaxSparseKVPool(
size=8,
page_size=1,
dtype=torch.float32,
head_num=1,
head_dim=8,
idx_head_dim=16,
dense_layer_ids=[0],
sparse_layer_ids=[1],
disable_value_sparse_layer_ids=[1],
device="cpu",
start_layer=0,
end_layer=2,
enable_hisparse=True,
)
allocator = HiSparseTokenToKVPoolAllocator(
size=pool.size,
page_size=1,
dtype=pool.dtype,
device="cpu",
kvcache=pool,
need_sort=False,
)
coordinator = HiSparseCoordinator.__new__(HiSparseCoordinator)
coordinator.is_dsv4_hisparse = False
coordinator.mem_pool_device = pool.main_pool
coordinator.token_to_kv_pool_allocator = allocator
coordinator.device_buffer_size = 2
coordinator.req_to_device_buffer = allocator.hisparse_attn_allocator.alloc(
3
).reshape(1, 3)
coordinator.req_device_buffer_size = torch.tensor([3])
coordinator.req_device_buffer_token_locs = torch.zeros(
(1, 1, 3), dtype=torch.int32
)
coordinator._skip_first_backup = [True]
out_loc = allocator.alloc(1)
with patch("sglang.srt.managers.hisparse_coordinator._is_hip", False):
for _ in range(2):
coordinator._skip_first_backup[0] = True
coordinator.map_last_loc_to_buffer(
seq_lens=torch.tensor([3]),
out_cache_loc=out_loc,
req_pool_indices=torch.tensor([0]),
seq_lens_cpu=torch.tensor([3]),
req_pool_indices_cpu=torch.tensor([0]),
)
self.assertEqual(
allocator.hisparse_attn_allocator.available_size(), pool.size - 3
)
torch.testing.assert_close(
allocator.full_to_hisparse_device_index_mapping[out_loc],
coordinator.req_to_device_buffer[:, 2],
)
class TestDeepSeekV4HiSparseAllocator(CustomTestCase):
def setUp(self):
# The code under test reads its config from the bags.
@@ -1,14 +1,23 @@
import unittest
from unittest.mock import patch
import torch
from sglang.srt.disaggregation.utils import get_kv_transfer_buf_infos
from sglang.srt.mem_cache.memory_pool import MiniMaxSparseKVPool
from sglang.srt.mem_cache.pool_host.mha import (
HiSparseMHATokenToKVPoolHost,
MHATokenToKVPoolHost,
)
from sglang.test.ci.ci_register import register_cpu_ci
from sglang.test.test_utils import CustomTestCase
register_cpu_ci(est_time=10, suite="base-a-test-cpu")
def _make_k_only_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
def _make_k_only_pool(
start_layer: int = 0, *, enable_hisparse: bool = False
) -> MiniMaxSparseKVPool:
"""Mirror the released MiniMax-M3 config shape: all sparse layers K-only."""
dense_layer_ids = [start_layer, start_layer + 1, start_layer + 2]
sparse_layer_ids = [start_layer + 3 + i for i in range(4)]
@@ -26,10 +35,11 @@ def _make_k_only_pool(start_layer: int = 0) -> MiniMaxSparseKVPool:
device="cpu",
start_layer=start_layer,
end_layer=end_layer,
enable_hisparse=enable_hisparse,
)
class TestMiniMaxSparsePoolPD(unittest.TestCase):
class TestMiniMaxSparsePoolPD(CustomTestCase):
def test_contiguous_buf_infos_main_only(self):
pool = _make_k_only_pool()
ptrs, lens, item_lens = pool.get_contiguous_buf_infos()
@@ -53,6 +63,52 @@ class TestMiniMaxSparsePoolPD(unittest.TestCase):
self.assertEqual(lens[i], buf.nbytes)
self.assertEqual(item_lens[i], buf[0].nbytes * pool.page_size)
def test_hisparse_host_registration(self):
"""PD startup must register every sparse host K/V buffer with page strides."""
pool = _make_k_only_pool(enable_hisparse=True)
host = HiSparseMHATokenToKVPoolHost.__new__(HiSparseMHATokenToKVPoolHost)
with patch(
"sglang.srt.mem_cache.pool_host.base.host_memory_budget_bytes",
return_value=1 << 30,
):
MHATokenToKVPoolHost.__init__(
host,
device_pool=pool.main_pool,
host_to_device_ratio=2,
host_size=0,
page_size=pool.page_size,
layout="layer_first",
pin_memory=False,
)
ptrs, lens, item_lens = get_kv_transfer_buf_infos(host)
buffers = list(host.k_buffer.unbind()) + list(host.v_buffer.unbind())
self.assertEqual(len(buffers), 8)
self.assertEqual(ptrs, [buffer.data_ptr() for buffer in buffers])
self.assertEqual(lens, [buffer.nbytes for buffer in buffers])
self.assertEqual(
item_lens, [buffer[0].nbytes * pool.page_size for buffer in buffers]
)
def test_pd_registration_separates_dense_and_sparse_layers(self):
"""Both PD peers must keep dense device KV out of the sparse transfer list."""
for hisparse in (False, True):
pool = _make_k_only_pool(enable_hisparse=hisparse)
for layers, infos in (
(range(3, 7), get_kv_transfer_buf_infos(pool)),
(range(3), pool.get_dense_kv_state_buf_infos()),
):
with self.subTest(hisparse=hisparse, layers=layers):
buffers = [pool.get_key_buffer(i) for i in layers] + [
pool.get_value_buffer(i) for i in layers
]
ptrs, lens, item_lens = infos
self.assertEqual(ptrs, [buffer.data_ptr() for buffer in buffers])
self.assertEqual(lens, [buffer.nbytes for buffer in buffers])
self.assertEqual(
item_lens,
[buffer[0].nbytes * pool.page_size for buffer in buffers],
)
if __name__ == "__main__":
unittest.main()