diff --git a/python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh b/python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh index 6dcf77e9b943..53c733b45168 100644 --- a/python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh +++ b/python/sglang/kernels/jit/csrc/attention/qsa_indexer.cuh @@ -277,11 +277,14 @@ struct QsaIndexKCompressParams { const float* cos_sin_cache; // [positions_capacity, rotary_dim] const int32_t* axis_map; // [rotary_dim / 2] const void* weight; // [kHeadDim] - const int32_t* write_locs; // [groups] + const int32_t* owner_mask; // [groups] + const int32_t* local_write_locs; // [groups] void* compressed_k_buffer; // [compressed_slots, kHeadDim] int32_t compress_ratio; int32_t rotary_dim; int32_t num_groups; + int32_t state_slots; + int32_t compressed_slots; float eps; }; @@ -306,7 +309,24 @@ __launch_bounds__(128) void qsa_index_k_compress_kernel(const QsaIndexKCompressP device::PDLWaitPrimary(); + // Shape-derived eager launches retain padded and non-owner groups. Skip + // those before touching source state; their locations need not be valid. + if (params.owner_mask[group] == 0) { + device::PDLTriggerSecondary(); + return; + } const int32_t* locs = params.group_locs + group * params.compress_ratio; + const int32_t write_loc = params.local_write_locs[group]; + if (write_loc < 0 || write_loc >= params.compressed_slots) { + device::PDLTriggerSecondary(); + return; + } + for (int32_t r = 0; r < params.compress_ratio; ++r) { + if (locs[r] < 0 || locs[r] >= params.state_slots) { + device::PDLTriggerSecondary(); + return; + } + } const int32_t loc0 = locs[0]; // fp32 mean over the group, rounded to the storage dtype exactly like @@ -361,8 +381,7 @@ __launch_bounds__(128) void qsa_index_k_compress_kernel(const QsaIndexKCompressP pos[a] = params.rope_position_buffer[static_cast(loc0) * 3 + a]; } - TOut* out_row = - static_cast(params.compressed_k_buffer) + static_cast(params.write_locs[group]) * kHeadDim; + TOut* out_row = static_cast(params.compressed_k_buffer) + static_cast(write_loc) * kHeadDim; qsa_mrope_apply( smem_rows[warp], out_row, params.cos_sin_cache, params.axis_map, pos, params.rotary_dim); @@ -453,7 +472,8 @@ void qsa_index_k_compress( tvm::ffi::TensorView cos_sin_cache, tvm::ffi::TensorView axis_map, tvm::ffi::TensorView weight, - tvm::ffi::TensorView write_locs, + tvm::ffi::TensorView owner_mask, + tvm::ffi::TensorView local_write_locs, tvm::ffi::TensorView compressed_k_buffer, int64_t compress_ratio, int64_t rotary_dim, @@ -472,7 +492,8 @@ void qsa_index_k_compress( TensorMatcher({cache_rows, rotary_dim}).with_dtype().with_device(device).verify(cos_sin_cache); TensorMatcher({rotary_dim / 2}).with_dtype().with_device(device).verify(axis_map); TensorMatcher({D}).with_dtype().with_device(device).verify(weight); - TensorMatcher({groups}).with_dtype().with_device(device).verify(write_locs); + TensorMatcher({groups}).with_dtype().with_device(device).verify(owner_mask); + TensorMatcher({groups}).with_dtype().with_device(device).verify(local_write_locs); auto compressed_slots = SymbolicSize{"compressed_slots"}; TensorMatcher({compressed_slots, D}).with_dtype().with_device(device).verify(compressed_k_buffer); @@ -490,11 +511,14 @@ void qsa_index_k_compress( .cos_sin_cache = static_cast(cos_sin_cache.data_ptr()), .axis_map = static_cast(axis_map.data_ptr()), .weight = weight.data_ptr(), - .write_locs = static_cast(write_locs.data_ptr()), + .owner_mask = static_cast(owner_mask.data_ptr()), + .local_write_locs = static_cast(local_write_locs.data_ptr()), .compressed_k_buffer = compressed_k_buffer.data_ptr(), .compress_ratio = static_cast(compress_ratio), .rotary_dim = static_cast(rotary_dim), .num_groups = static_cast(num_groups), + .state_slots = static_cast(slots.unwrap()), + .compressed_slots = static_cast(compressed_slots.unwrap()), .eps = eps, }; LaunchKernel(static_cast(div_ceil(num_groups, 4)), 128, device.unwrap()) diff --git a/python/sglang/kernels/ops/attention/dcp_kernels.py b/python/sglang/kernels/ops/attention/dcp_kernels.py index 5842702cf08a..db4b79d44390 100644 --- a/python/sglang/kernels/ops/attention/dcp_kernels.py +++ b/python/sglang/kernels/ops/attention/dcp_kernels.py @@ -603,6 +603,8 @@ def _dcp_lse_combine_kernel( out_stride_B, out_stride_H, out_stride_D, + out_lse_stride_B, + out_lse_stride_H, N: tl.constexpr, HEAD_DIM: tl.constexpr, IS_BASE_E: tl.constexpr, @@ -635,7 +637,8 @@ def _dcp_lse_combine_kernel( ) lse_max = tl.where(lse_i > lse_max, lse_i, lse_max) - lse_max = tl.where(lse_max == -float("inf"), 0.0, lse_max) + has_any = lse_max != -float("inf") + lse_max = tl.where(has_any, lse_max, 0.0) # Pass 2: accumulate weighted outputs weight_sum = tl.zeros([], dtype=tl.float32) @@ -651,6 +654,7 @@ def _dcp_lse_combine_kernel( w = tl.exp(centered) else: w = tl.exp2(centered) + w = tl.where(lse_i == -float("inf"), 0.0, w) weight_sum += w o_offsets = ( @@ -662,7 +666,7 @@ def _dcp_lse_combine_kernel( partial_out = tl.load(recv_output_ptr + o_offsets).to(tl.float32) acc += partial_out * w - acc = acc / tl.maximum(weight_sum, 1.0) + acc = tl.where(has_any, acc / tl.maximum(weight_sum, 1.0), 0.0) out_offsets = ( batch_idx * out_stride_B + head_idx * out_stride_H + d_offsets * out_stride_D @@ -671,10 +675,10 @@ def _dcp_lse_combine_kernel( if RETURN_LSE: if IS_BASE_E: - global_lse = tl.log(weight_sum) + lse_max + global_lse = tl.where(has_any, tl.log(weight_sum) + lse_max, -float("inf")) else: - global_lse = tl.log2(weight_sum) + lse_max - out_lse_offset = batch_idx * recv_lse_stride_B + head_idx * recv_lse_stride_H + global_lse = tl.where(has_any, tl.log2(weight_sum) + lse_max, -float("inf")) + out_lse_offset = batch_idx * out_lse_stride_B + head_idx * out_lse_stride_H tl.store(out_lse_ptr + out_lse_offset, global_lse) @@ -722,6 +726,8 @@ def dcp_lse_combine_triton( out.stride(0), out.stride(1), out.stride(2), + out_lse.stride(0) if return_lse else 0, + out_lse.stride(1) if return_lse else 0, N=N, HEAD_DIM=D, IS_BASE_E=is_lse_base_on_e, diff --git a/python/sglang/kernels/ops/attention/qsa_indexer.py b/python/sglang/kernels/ops/attention/qsa_indexer.py index dcfa1c5ecbaf..c327b84c5284 100644 --- a/python/sglang/kernels/ops/attention/qsa_indexer.py +++ b/python/sglang/kernels/ops/attention/qsa_indexer.py @@ -24,6 +24,8 @@ from typing import TYPE_CHECKING, Optional import torch +import triton +import triton.language as tl from sglang.kernels.jit.utils import ( cache_once, @@ -36,6 +38,49 @@ from tvm_ffi.module import Module +@triton.jit +def _qsa_store_compressed_keys_kernel( + source, + cache, + local_locs, + owner_mask, + WIDTH: tl.constexpr, + BLOCK: tl.constexpr, +): + row = tl.program_id(0) + if tl.load(owner_mask + row): + slot = tl.load(local_locs + row).to(tl.int64) + offsets = tl.arange(0, BLOCK) + values = tl.load(source + row * WIDTH + offsets, mask=offsets < WIDTH) + tl.store(cache + slot * WIDTH + offsets, values, mask=offsets < WIDTH) + + +def qsa_store_compressed_keys( + compressed_keys: torch.Tensor, + cache: torch.Tensor, + local_locs: torch.Tensor, + owner_mask: torch.Tensor, +) -> None: + """Store graph fallback keys at local slots without dynamic owner filtering. + + Active destinations must be unique. Masked rows may share a safe slot, + including slot zero even when another row owns that real local slot. + """ + rows = compressed_keys.shape[0] + if rows == 0: + return + keys = compressed_keys.flatten(1).contiguous().to(cache.dtype) + width = keys.shape[1] + _qsa_store_compressed_keys_kernel[(rows,)]( + keys, + cache, + local_locs, + owner_mask, + WIDTH=width, + BLOCK=triton.next_power_of_2(width), + ) + + @cache_once def _jit_qsa_indexer_module( dtype: torch.dtype, out_dtype: torch.dtype, head_dim: int, is_neox_style: bool @@ -147,7 +192,8 @@ def qsa_index_k_compress_store( cos_sin_cache: torch.Tensor, axis_map: torch.Tensor, weight: torch.Tensor, - write_locs: torch.Tensor, + owner_mask: torch.Tensor, + local_write_locs: torch.Tensor, compressed_k_buffer: torch.Tensor, compress_ratio: int, rotary_dim: int, @@ -167,7 +213,9 @@ def qsa_index_k_compress_store( cos_sin_cache : CUDA fp32 [capacity, rotary_dim] RoPE cache axis_map : CUDA int32 [rotary_dim // 2] position-axis per pair index weight : [head_dim] gemma norm weight (kernel applies 1 + w) - write_locs : CUDA int32 [groups] compressed-cache slots to write + owner_mask : CUDA int32 [groups], zero skips all reads and writes + local_write_locs : CUDA int32 [groups] rank-local compressed-cache slots; + ignored where owner_mask is zero compressed_k_buffer : CUDA [compressed_slots, head_dim] (written), ring dtype or fp8 e4m3 compress_ratio : raw keys per compressed key @@ -186,7 +234,8 @@ def qsa_index_k_compress_store( cos_sin_cache, axis_map, weight, - write_locs, + owner_mask, + local_write_locs, compressed_k_buffer, compress_ratio, rotary_dim, diff --git a/python/sglang/srt/arg_groups/fields/parallel.py b/python/sglang/srt/arg_groups/fields/parallel.py index 38ea4438bb54..d0ad3caa2452 100644 --- a/python/sglang/srt/arg_groups/fields/parallel.py +++ b/python/sglang/srt/arg_groups/fields/parallel.py @@ -60,6 +60,15 @@ class Parallel(msgspec.Struct): aliases=["--decode-context-parallel-size"], ), ] = 1 + qsa_cache_sharding_size: A[ + int, + Arg( + help="The number of tensor-parallel ranks over which QSA cache, " + "distributed Top-K, and owner attention are sharded. GDN, PLE, and " + "the QSA pending ring remain replicated.", + aliases=["--state-cp-size"], + ), + ] = 1 pp_size: A[ int, Arg( @@ -391,6 +400,9 @@ class Parallel(msgspec.Struct): "zero where decode context parallelism is off." ) ) + qsa_cache_sharding_rank = Derived( + doc="This process's rank in its QSA cache-sharding group." + ) attn_dp_rank = Derived( doc=( "This process's index in the attention-DP group, computed from " @@ -431,3 +443,4 @@ class Parallel(msgspec.Struct): attn_cp_group = Derived(doc="The attention context-parallel group.") shared_experts_tp_group = Derived(doc=("The shared-expert tensor-parallel group.")) dcp_group = Derived(doc="The decode context-parallel group.") + qsa_cache_sharding_group = Derived(doc="The QSA cache-sharding group.") diff --git a/python/sglang/srt/arg_groups/parallel_hook.py b/python/sglang/srt/arg_groups/parallel_hook.py index 21fa2dbb7b18..f5810297bdc2 100644 --- a/python/sglang/srt/arg_groups/parallel_hook.py +++ b/python/sglang/srt/arg_groups/parallel_hook.py @@ -285,6 +285,38 @@ def handle_decode_context_parallelism(server_args: Any): f"got --dcp-comm-backend={cfg.dcp_comm_backend}." ) + from sglang.srt.layers.attention.qsa.cache_sharding import ( + validate_qsa_cache_sharding, + ) + + if cfg.qsa_cache_sharding_size == 1: + return + if parse_connector_type(cfg.model_path) == ConnectorType.INSTANCE: + raise ValueError("QSA cache sharding requires a local model configuration") + model_config = model_config_of(server_args) + hf_config = model_config.hf_config + from sglang.srt.layers.attention.qsa.config import is_qwen_qsa + + validate_qsa_cache_sharding( + device=cfg.device, + is_qsa=is_qwen_qsa(hf_config), + tp_size=derive_attn_tp_size( + tp_size=cfg.tp_size, + attn_cp_size=resolved_view(server_args).attn_cp_size, + attn_dp_size=cfg.attn_dp_size, + ), + sharding_size=cfg.qsa_cache_sharding_size, + num_attention_heads=model_config.num_attention_heads, + num_key_value_heads=model_config.num_key_value_heads, + pp_size=cfg.pp_size, + attn_cp_size=cfg.attn_cp_size, + enable_prefill_cp=cfg.enable_prefill_cp, + ep_size=cfg.ep_size, + moe_dp_size=cfg.moe_dp_size, + disaggregation_mode=cfg.disaggregation_mode, + dcp_size=cfg.dcp_size, + ) + def handle_data_parallelism(server_args: Any): # The resets without attention DP live in the resolution pipeline diff --git a/python/sglang/srt/distributed/parallel_state.py b/python/sglang/srt/distributed/parallel_state.py index 362e2a5d3a40..5824e9475203 100644 --- a/python/sglang/srt/distributed/parallel_state.py +++ b/python/sglang/srt/distributed/parallel_state.py @@ -2234,6 +2234,7 @@ def init_model_parallel_group( _SHARED_EXPERTS_TP: Optional[GroupCoordinator] = None _ATTN_CP: Optional[GroupCoordinator] = None _DCP: Optional[GroupCoordinator] = None +_QSA_CACHE_SHARDING: Optional[GroupCoordinator] = None def get_tp_group() -> GroupCoordinator: @@ -2260,6 +2261,10 @@ def get_dcp_group() -> GroupCoordinator: return get_parallel().dcp_group +def get_qsa_cache_sharding_group() -> GroupCoordinator: + return get_parallel().qsa_cache_sharding_group + + _MOE_DP: Optional[GroupCoordinator] = None _MOE_EP: Optional[GroupCoordinator] = None _MOE_TP: Optional[GroupCoordinator] = None @@ -2330,7 +2335,14 @@ def graph_capture(stream=None): ): with contextlib.ExitStack() as stack: seen = {id(_TP), id(_PP)} - for group in (_DCP, _ATTN_TP, _SHARED_EXPERTS_TP, _MOE_EP, _MOE_TP): + for group in ( + _DCP, + _QSA_CACHE_SHARDING, + _ATTN_TP, + _SHARED_EXPERTS_TP, + _MOE_EP, + _MOE_TP, + ): if group is not None and id(group) not in seen: seen.add(id(group)) stack.enter_context(group.graph_capture(context)) @@ -2574,6 +2586,16 @@ def init_distributed_environment( get_parallel().override_permanently(world_group=_WORLD) +def _build_tp_subgroups( + tp_group_ranks: list[list[int]], subgroup_size: int +) -> list[list[int]]: + return [ + tp_group[start : start + subgroup_size] + for tp_group in tp_group_ranks + for start in range(0, len(tp_group), subgroup_size) + ] + + def initialize_model_parallel( backend: Optional[str] = None, enable_symm_mem: bool = False, @@ -2644,6 +2666,7 @@ def initialize_model_parallel( attention_context_model_parallel_size = parallel.attn_cp_size moe_data_model_parallel_size = parallel.moe_dp_size decode_context_parallel_size = parallel.attn_dcp_size + qsa_cache_sharding_size = parallel.qsa_cache_sharding_size shared_experts_tensor_parallel_size = parallel.shared_experts_tp_size # Joiners construct their local TP/PP layout in global rank space. @@ -2677,6 +2700,15 @@ def initialize_model_parallel( f"tensor_model_parallel_size ({tensor_model_parallel_size}) must be divisible by " f"decode_context_parallel_size ({decode_context_parallel_size})" ) + if qsa_cache_sharding_size < 1: + raise RuntimeError( + f"qsa_cache_sharding_size ({qsa_cache_sharding_size}) must be >= 1" + ) + if tensor_model_parallel_size % qsa_cache_sharding_size != 0: + raise RuntimeError( + f"tensor_model_parallel_size ({tensor_model_parallel_size}) must be divisible by " + f"qsa_cache_sharding_size ({qsa_cache_sharding_size})" + ) # Build the tensor model-parallel groups. num_tensor_model_parallel_groups: int = world_size // tensor_model_parallel_size @@ -2708,14 +2740,8 @@ def initialize_model_parallel( global _DCP assert _DCP is None, "decode context parallel group is already initialized" if decode_context_parallel_size > 1: - dcp_group_ranks = [] - for tp_group in group_ranks: - for start in range(0, len(tp_group), decode_context_parallel_size): - dcp_group_ranks.append( - tp_group[start : start + decode_context_parallel_size] - ) _DCP = init_model_parallel_group( - dcp_group_ranks, + _build_tp_subgroups(group_ranks, decode_context_parallel_size), get_world_group().local_rank, backend, # Current DCP paths use all-gather/reduce-scatter, all-to-all, or @@ -2733,6 +2759,29 @@ def initialize_model_parallel( f"DCP enabled, dcp_size={decode_context_parallel_size}, tp_size={tensor_model_parallel_size}" ) + # QSA cache sharding is independent of decode context parallelism. + global _QSA_CACHE_SHARDING + assert _QSA_CACHE_SHARDING is None, ( + "QSA cache-sharding group is already initialized" + ) + if qsa_cache_sharding_size > 1: + _QSA_CACHE_SHARDING = init_model_parallel_group( + _build_tp_subgroups(group_ranks, qsa_cache_sharding_size), + get_world_group().local_rank, + backend, + use_message_queue_broadcaster=envs.SGLANG_USE_MESSAGE_QUEUE_BROADCASTER.get(), + group_name="qsa_cache_sharding", + recovered_rank=recovered_rank, + rank_offset=rank_offset, + max_world_size=max_world_size, + ) + if _TP.rank_in_group == 0: + logger.info( + "QSA cache sharding enabled, size=%d, tp_size=%d", + qsa_cache_sharding_size, + tensor_model_parallel_size, + ) + attn_dp_size = attention_data_parallel_size attn_cp_size = attention_context_model_parallel_size # The groups below are built at these numbers, and the same dict is stamped @@ -2992,6 +3041,7 @@ def initialize_model_parallel( "attn_cp_group": _ATTN_CP, "shared_experts_tp_group": _SHARED_EXPERTS_TP, "dcp_group": _DCP, + "qsa_cache_sharding_group": _QSA_CACHE_SHARDING, } get_parallel().override_permanently( **{name: group for name, group in built.items() if group is not None} @@ -3294,6 +3344,11 @@ def destroy_model_parallel(): _DCP.destroy() _DCP = None + global _QSA_CACHE_SHARDING + if _QSA_CACHE_SHARDING: + _QSA_CACHE_SHARDING.destroy() + _QSA_CACHE_SHARDING = None + global _MOE_EP if _MOE_EP: _MOE_EP.destroy() @@ -3488,6 +3543,7 @@ def monkey_patch_vllm_parallel_state(reverse: bool = False): "get_attn_tp_group": "attn_tp_group", "get_attn_cp_group": "attn_cp_group", "get_shared_experts_tp_group": "shared_experts_tp_group", + "get_qsa_cache_sharding_group": "qsa_cache_sharding_group", "get_dcp_group": "dcp_group", "get_world_size": "launch_world_size", "get_world_rank": "launch_world_rank", diff --git a/python/sglang/srt/layers/attention/qsa/cache_sharding.py b/python/sglang/srt/layers/attention/qsa/cache_sharding.py new file mode 100644 index 000000000000..3f4baa7e762b --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/cache_sharding.py @@ -0,0 +1,203 @@ +from __future__ import annotations + +import os +from dataclasses import dataclass + +import torch +import torch.distributed as dist + + +@dataclass(frozen=True) +class QSACacheShardingRuntime: + size: int + rank: int + group: object | None + + @property + def enabled(self) -> bool: + return self.size > 1 + + +def resolve_qsa_cache_sharding_runtime(parallel) -> QSACacheShardingRuntime: + size = int(getattr(parallel, "qsa_cache_sharding_size", 1)) + return QSACacheShardingRuntime( + size=size, + rank=int(getattr(parallel, "qsa_cache_sharding_rank", 0)), + group=getattr(parallel, "qsa_cache_sharding_group", None) if size > 1 else None, + ) + + +def get_qsa_cache_sharding_runtime(parallel=None) -> QSACacheShardingRuntime: + from sglang.srt import runtime_context + + if parallel is None: + parallel = runtime_context.get_parallel() + try: + return resolve_qsa_cache_sharding_runtime(parallel) + except ValueError as exc: + if "not published" not in str(exc): + raise + return QSACacheShardingRuntime(size=1, rank=0, group=None) + + +def assert_qsa_cache_sharding_runtime_match( + expected: QSACacheShardingRuntime, + actual: QSACacheShardingRuntime, + *, + component: str, +) -> None: + expected_topology = (expected.size, expected.rank) + actual_topology = (actual.size, actual.rank) + if expected_topology != actual_topology: + raise ValueError( + f"{component} QSA cache-sharding runtime mismatch: " + f"size/rank={actual_topology!r} (expected {expected_topology!r})" + ) + if expected.enabled and expected.group is not actual.group: + raise ValueError(f"{component} QSA cache-sharding group does not match") + + +def validate_qsa_cache_sharding( + *, + device: str, + is_qsa: bool, + tp_size: int, + sharding_size: int, + num_attention_heads: int, + num_key_value_heads: int, + pp_size: int, + attn_cp_size: int, + enable_prefill_cp: bool, + ep_size: int, + moe_dp_size: int, + disaggregation_mode: str, + dcp_size: int = 1, +) -> None: + if sharding_size == 1: + return + if device != "cuda": + raise ValueError("QSA cache sharding currently requires CUDA") + if not is_qsa: + raise ValueError("QSA cache sharding requires a QSA model") + if sharding_size < 1 or tp_size % sharding_size: + raise ValueError("QSA cache sharding requires its size to divide tp_size") + if ( + min(num_attention_heads, num_key_value_heads) < 1 + or num_attention_heads % tp_size + or num_attention_heads % num_key_value_heads + or tp_size % num_key_value_heads + or (tp_size // num_key_value_heads) % sharding_size + ): + raise ValueError( + "QSA cache sharding must divide the KV-head replica group size: " + f"attention_tp_size={tp_size}, num_attention_heads={num_attention_heads}, " + f"num_key_value_heads={num_key_value_heads}, sharding_size={sharding_size}" + ) + unsupported = ( + (dcp_size != 1, "decode context parallelism"), + (pp_size != 1, "pipeline parallelism"), + (attn_cp_size != 1, "attention context parallelism"), + (enable_prefill_cp, "prefill CP"), + (disaggregation_mode != "null", "PD disaggregation"), + ) + for enabled, feature in unsupported: + if enabled: + raise ValueError(f"QSA cache sharding does not support {feature}") + if ep_size < 1 or moe_dp_size < 1 or tp_size % (ep_size * moe_dp_size): + raise ValueError("QSA cache sharding requires a valid MoE EP topology") + + +class QSACacheShardingStaticBuffers: + def __init__( + self, + *, + max_rows: int, + topk: int, + owner_topk: int | None = None, + num_heads: int, + head_dim: int, + world_size: int = 1, + device: torch.device | str, + dtype: torch.dtype, + candidate_group=None, + ) -> None: + owner_topk = topk if owner_topk is None else owner_topk + if min(max_rows, topk, owner_topk, num_heads, head_dim, world_size) <= 0: + raise ValueError("static buffer dimensions must be positive") + if num_heads % world_size: + raise ValueError("num_heads must be divisible by world_size") + self.max_rows, self.world_size = max_rows, world_size + self.topk, self.owner_topk = topk, owner_topk + self.num_heads, self.head_dim = num_heads, head_dim + self._device, self._dtype = torch.device(device), dtype + local_heads = num_heads // world_size + + def empty(shape, storage_dtype=dtype): + return torch.empty(shape, dtype=storage_dtype, device=device) + + self.candidate_transport = empty( + (max_rows, topk, 2), storage_dtype=torch.float32 + ) + self.candidate_recv = empty((world_size, max_rows, topk, 2), torch.float32) + self.candidate_symm_state = self._create_candidate_symm_state(candidate_group) + self.candidate_scores = empty((max_rows, topk), torch.float32) + self.candidate_indices = empty((max_rows, topk), torch.int32) + self.q_send = empty((local_heads, max_rows, head_dim)) + self.q_recv = empty((num_heads, max_rows, head_dim)) + self.q_gathered = empty((max_rows, num_heads, head_dim)) + self.local_slots = empty((max_rows, owner_topk), torch.int64) + self.owner_lengths = empty((max_rows,), torch.int32) + self.owner_a2a_send = empty( + (world_size, max_rows, local_heads, head_dim + 1), torch.float32 + ) + self.owner_a2a_recv = torch.empty_like(self.owner_a2a_send) + + def _create_candidate_symm_state(self, candidate_group): + if candidate_group is None: + return None + from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( + create_state, + ) + + device_group = getattr(candidate_group, "device_group", None) + capability = ( + torch.cuda.get_device_capability(self._device) + if self._device.type == "cuda" + else (0, 0) + ) + local_supported = ( + os.environ.get("SGLANG_QSA_CANDIDATE_SYMM_MEM", "1") == "1" + and self.world_size == 4 + and self._device.type == "cuda" + and capability[0] == 9 + and self.topk * 4 % 8 == 0 + and self.candidate_transport.data_ptr() % 16 == 0 + and type(device_group) is dist.ProcessGroup + ) + if device_group is None: + return None + consensus = torch.tensor( + int(local_supported), dtype=torch.int32, device=self._device + ) + dist.all_reduce(consensus, op=dist.ReduceOp.MIN, group=device_group) + if not bool(consensus.item()): + return None + state = create_state( + group=device_group, + rank_in_group=int(candidate_group.rank_in_group), + max_tokens=self.max_rows, + hidden_size=self.topk * 4 * self.world_size, + device=self._device, + ) + multicast = torch.tensor( + int(state.symm_mem_hdl.multicast_ptr != 0), + dtype=torch.int32, + device=self._device, + ) + dist.all_reduce(multicast, op=dist.ReduceOp.MIN, group=device_group) + return state if bool(multicast.item()) else None + + def for_batch(self, rows: int) -> QSACacheShardingStaticBuffers: + if rows < 0 or rows > self.max_rows: + raise ValueError(f"batch rows {rows} exceed capacity {self.max_rows}") + return self diff --git a/python/sglang/srt/layers/attention/qsa/distributed_topk.py b/python/sglang/srt/layers/attention/qsa/distributed_topk.py new file mode 100644 index 000000000000..9b17bec80042 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/distributed_topk.py @@ -0,0 +1,355 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from sglang.srt.distributed.device_communicators.triton_symm_mem_ag import ( + _blockwise_barrier, + _multimem_st_128, + _sync_threads, + all_gather_inner, +) +from sglang.srt.distributed.utils import all_gather_single + + +def all_gather_qsa_topk_candidates( + candidate_scores: torch.Tensor, + candidate_global_indices: torch.Tensor, + *, + group, + static_buffers=None, +) -> torch.Tensor: + if candidate_global_indices.shape != candidate_scores.shape: + raise ValueError("candidate scores and indices must have matching shapes") + if candidate_global_indices.dtype not in (torch.int32, torch.int64): + raise TypeError("candidate indices must have integer dtype") + if static_buffers is None: + packed = torch.empty( + (*candidate_scores.shape, 2), + dtype=torch.float32, + device=candidate_scores.device, + ) + packed[..., 0].copy_(candidate_scores) + packed[..., 1].view(torch.int32).copy_(candidate_global_indices) + else: + rows, width = candidate_scores.shape + static_buffers = static_buffers.for_batch(rows) + packed = static_buffers.candidate_transport + packed[..., 0].fill_(float("-inf")) + packed[..., 1].view(torch.int32).fill_(-1) + packed[:rows, :width, 0].copy_(candidate_scores) + packed[:rows, :width, 1].view(torch.int32).copy_(candidate_global_indices) + with torch.profiler.record_function("qsa.candidate_gather"): + if static_buffers is None: + return group.all_gather(packed, dim=1) + symm_state = getattr(static_buffers, "candidate_symm_state", None) + if symm_state is not None: + packed_bf16 = packed.view(torch.bfloat16).reshape( + static_buffers.max_rows, -1 + ) + gathered_bf16 = all_gather_inner( + symm_state, + packed_bf16, + tp_hidden_dim=packed_bf16.shape[1] * group.world_size, + safe=False, + skip_entry_sync=True, + ) + return gathered_bf16.view(torch.float32).reshape( + static_buffers.max_rows, group.world_size * width, 2 + )[:rows] + recv = static_buffers.candidate_recv.view( + group.world_size * static_buffers.max_rows, + static_buffers.topk, + 2, + ) + if not packed.is_cuda: + group.all_gather_into_tensor(recv, packed) + else: + device_group = getattr(group, "device_group", None) + if device_group is None: + group.all_gather_into_tensor(recv, packed) + else: + work = all_gather_single( + recv, + packed, + group=device_group, + async_op=True, + ) + work.block_current_stream() + return static_buffers.candidate_recv[:, :rows, :width] + + +@triton.jit +def _fused_publish_qsa_topk_candidates_kernel( + logits_ptr, + local_indices_ptr, + row_starts_ptr, + logical_positions_ptr, + multicast_ptr, + signal_pad_ptr, + logits_row_stride, + logical_row_stride, + rows, + TOPK: tl.constexpr, + RANK: tl.constexpr, + WORLD_SIZE: tl.constexpr, + BLOCK_SIZE: tl.constexpr, +): + # A rank may reach this layer in the next replay while a peer is still + # merging the previous payload. Synchronize before overwriting the shared + # fixed-address transport, then use the exit barrier for publication. + _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem="relaxed") + _sync_threads() + + # One 128-bit multimem store publishes two {fp32 score, int32 id} records. + chunks_per_row: tl.constexpr = TOPK // 2 + total_chunks = rows * chunks_per_row + block_start = tl.program_id(0) * BLOCK_SIZE + grid_stride: tl.constexpr = 4 * BLOCK_SIZE + while block_start < total_chunks: + chunk = block_start + tl.arange(0, BLOCK_SIZE) + valid_chunk = chunk < total_chunks + row = chunk // chunks_per_row + first_col = (chunk % chunks_per_row) * 2 + first_offset = row * TOPK + first_col + first_local = tl.load( + local_indices_ptr + first_offset, mask=valid_chunk, other=-1 + ).to(tl.int32) + second_local = tl.load( + local_indices_ptr + first_offset + 1, mask=valid_chunk, other=-1 + ).to(tl.int32) + row_start = tl.load(row_starts_ptr + row, mask=valid_chunk, other=0) + first_valid = valid_chunk & (first_local >= 0) + second_valid = valid_chunk & (second_local >= 0) + first_score = tl.load( + logits_ptr + row * logits_row_stride + row_start + first_local, + mask=first_valid, + other=-float("inf"), + ).to(tl.float32) + second_score = tl.load( + logits_ptr + row * logits_row_stride + row_start + second_local, + mask=second_valid, + other=-float("inf"), + ).to(tl.float32) + first_id = tl.load( + logical_positions_ptr + row * logical_row_stride + first_local, + mask=first_valid, + other=-1, + ).to(tl.int32) + second_id = tl.load( + logical_positions_ptr + row * logical_row_stride + second_local, + mask=second_valid, + other=-1, + ).to(tl.int32) + first_id = tl.where(first_valid, first_id, -1) + second_id = tl.where(second_valid, second_id, -1) + + # The symmetric buffer is bf16 storage viewed as uint32 words. Each + # rank owns TOPK*2 words per row; each chunk writes four words. + out_word = row * WORLD_SIZE * TOPK * 2 + RANK * TOPK * 2 + first_col * 2 + out_ptr = multicast_ptr.to(tl.int64).to(tl.pointer_type(tl.uint32)) + out_word + _multimem_st_128( + out_ptr, + first_score.to(tl.uint32, bitcast=True), + first_id.to(tl.uint32, bitcast=True), + second_score.to(tl.uint32, bitcast=True), + second_id.to(tl.uint32, bitcast=True), + valid_chunk, + ) + block_start += grid_stride + + _sync_threads() + _blockwise_barrier(signal_pad_ptr, RANK, WORLD_SIZE, sem="acq_rel") + + +def fused_publish_and_merge_qsa_topk_candidates( + logits: torch.Tensor, + local_indices: torch.Tensor, + row_starts: torch.Tensor, + logical_block_positions: torch.Tensor | None, + *, + group, + topk: int, + static_buffers, +) -> tuple[torch.Tensor, torch.Tensor]: + if logical_block_positions is None: + raise ValueError("fused candidate publish requires logical block positions") + rows = local_indices.shape[0] + static_buffers = static_buffers.for_batch(rows) + state = static_buffers.candidate_symm_state + if state is None: + raise RuntimeError("fused candidate publish requires symmetric memory") + if topk % 2: + raise ValueError("fused candidate publish requires an even topk") + if logits.dtype != torch.float32 or local_indices.dtype != torch.int32: + raise TypeError("fused candidate publish requires fp32 logits and int32 ids") + _fused_publish_qsa_topk_candidates_kernel[(4,)]( + logits, + local_indices, + row_starts, + logical_block_positions, + state.symm_mem_hdl.multicast_ptr, + state.symm_mem_hdl.signal_pad_ptrs_dev, + logits.stride(0), + logical_block_positions.stride(0), + rows, + TOPK=topk, + RANK=int(group.rank_in_group), + WORLD_SIZE=int(group.world_size), + BLOCK_SIZE=256, + num_warps=8, + ) + gathered = ( + state.comm_buff[:rows, : topk * 4 * group.world_size] + .view(torch.float32) + .reshape(rows, group.world_size * topk, 2) + ) + scores, indices = merge_qsa_topk_candidates(gathered, topk=topk) + static_buffers.candidate_scores[:rows, :topk].copy_(scores) + static_buffers.candidate_indices[:rows, :topk].copy_(indices) + return ( + static_buffers.candidate_scores[:rows, :topk], + static_buffers.candidate_indices[:rows, :topk], + ) + + +@triton.jit +def _merge_qsa_topk_candidates_kernel( + candidates_ptr, + output_scores_ptr, + output_indices_ptr, + width: tl.constexpr, + output_width: tl.constexpr, + block_width: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, block_width) + input_offsets = (row * width + offsets) * 2 + in_bounds = offsets < width + scores = tl.load( + candidates_ptr + input_offsets, mask=in_bounds, other=-float("inf") + ) + indices = tl.load( + candidates_ptr + input_offsets + 1, mask=in_bounds, other=0xFFFFFFFF + ).to(tl.int32, bitcast=True) + valid = in_bounds & (indices >= 0) & (scores == scores) # noqa: PLR0124 + + score_bits = scores.to(tl.int32, bitcast=True).to(tl.int64) & 0xFFFFFFFF + score_bits = tl.where(scores == 0.0, 0, score_bits) + ordered_score = tl.where( + (score_bits & 0x80000000) != 0, + (~score_bits) & 0xFFFFFFFF, + score_bits ^ 0x80000000, + ) + keys = (ordered_score - 0x80000000) * 0x100000000 + ( + 0xFFFFFFFF - indices.to(tl.int64) + ) + keys = tl.where(valid, keys, -0x8000000000000000) + keys = tl.sort(keys, dim=0, descending=True) + selected = offsets < output_width + selected_keys = tl.where(selected, keys, -0x8000000000000000) + selected_valid = selected & (selected_keys != -0x8000000000000000) + # Arithmetic shift preserves the exact high 32-bit score key. Integer + # division can round negative packed keys through fp64 and flip one ulp. + selected_ordered_score = (selected_keys >> 32) + 0x80000000 + selected_score_bits = tl.where( + (selected_ordered_score & 0x80000000) != 0, + selected_ordered_score ^ 0x80000000, + (~selected_ordered_score) & 0xFFFFFFFF, + ) + selected_scores = selected_score_bits.to(tl.int32).to(tl.float32, bitcast=True) + selected_ids = 0xFFFFFFFF - (selected_keys & 0xFFFFFFFF) + tl.store( + output_scores_ptr + row * output_width + offsets, + tl.where(selected_valid, selected_scores, -float("inf")), + mask=selected, + ) + tl.store( + output_indices_ptr + row * output_width + offsets, + tl.where(selected_valid, selected_ids, -1), + mask=selected, + ) + + +def _merge_qsa_topk_candidates_triton( + gathered_candidates: torch.Tensor, + *, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + rows, width, _ = gathered_candidates.shape + output_scores = torch.empty( + (rows, topk), dtype=torch.float32, device=gathered_candidates.device + ) + output_indices = torch.empty( + (rows, topk), dtype=torch.int32, device=gathered_candidates.device + ) + block_width = triton.next_power_of_2(width) + _merge_qsa_topk_candidates_kernel[(rows,)]( + gathered_candidates, + output_scores, + output_indices, + width=width, + output_width=topk, + block_width=block_width, + num_warps=4 if block_width <= 128 else 8, + ) + return output_scores, output_indices + + +def merge_qsa_topk_candidates( + gathered_candidates: torch.Tensor, + *, + topk: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if gathered_candidates.ndim == 4: + if gathered_candidates.shape[-1] != 2: + raise ValueError("candidate transport must end in score/index pairs") + gathered_candidates = ( + gathered_candidates.permute(1, 0, 2, 3) + .reshape(gathered_candidates.shape[1], -1, 2) + .contiguous() + ) + if gathered_candidates.ndim != 3 or gathered_candidates.shape[-1] != 2: + raise ValueError( + "gathered candidates must be [rows, width, 2] or " + f"[ranks, rows, width, 2], got {tuple(gathered_candidates.shape)}" + ) + if gathered_candidates.shape[1] < topk: + raise ValueError( + f"candidate width {gathered_candidates.shape[1]} is smaller than " + f"topk {topk}" + ) + if gathered_candidates.dtype != torch.float32: + raise TypeError("candidate transport must use float32") + if not gathered_candidates.is_cuda: + raise ValueError("distributed QSA candidate merge requires CUDA") + + with torch.profiler.record_function("qsa.candidate_merge"): + return _merge_qsa_topk_candidates_triton(gathered_candidates, topk=topk) + + +def gather_and_merge_qsa_topk_candidates( + candidate_scores: torch.Tensor, + candidate_global_indices: torch.Tensor, + *, + group, + topk: int, + static_buffers=None, +) -> tuple[torch.Tensor, torch.Tensor]: + gathered = all_gather_qsa_topk_candidates( + candidate_scores, + candidate_global_indices, + group=group, + static_buffers=static_buffers, + ) + scores, indices = merge_qsa_topk_candidates(gathered, topk=topk) + if static_buffers is None: + return scores, indices + rows = candidate_scores.shape[0] + static_buffers.candidate_scores[:rows, :topk].copy_(scores) + static_buffers.candidate_indices[:rows, :topk].copy_(indices) + return ( + static_buffers.candidate_scores[:rows, :topk], + static_buffers.candidate_indices[:rows, :topk], + ) diff --git a/python/sglang/srt/layers/attention/qsa/glue.py b/python/sglang/srt/layers/attention/qsa/glue.py index b824d8a72791..dd3facb11337 100644 --- a/python/sglang/srt/layers/attention/qsa/glue.py +++ b/python/sglang/srt/layers/attention/qsa/glue.py @@ -19,14 +19,19 @@ def build_qsa_indexer( raise ValueError( "build_qsa_indexer requires a config with a QSA indexer schema" ) + from sglang.srt.layers.attention.qsa.cache_sharding import ( + get_qsa_cache_sharding_runtime, + ) from sglang.srt.layers.attention.qsa.qsa_indexer import QSAIndexer + cache_sharding_runtime = get_qsa_cache_sharding_runtime() return QSAIndexer( config=config, layer_id=layer_id, quant_config=quant_config, prefix=prefix, rotary_emb=rotary_emb, + cache_sharding_runtime=cache_sharding_runtime, ) diff --git a/python/sglang/srt/layers/attention/qsa/graph_metadata.py b/python/sglang/srt/layers/attention/qsa/graph_metadata.py index 3ac793f96125..52aeaa360083 100644 --- a/python/sglang/srt/layers/attention/qsa/graph_metadata.py +++ b/python/sglang/srt/layers/attention/qsa/graph_metadata.py @@ -74,7 +74,9 @@ def _qsa_graph_row_metadata_kernel( # Graph output buffers. compressed_lens_ptr, write_locs_ptr, + write_owner_ptr, page_table_ptr, + block_positions_ptr, logical_positions_ptr, state_slots_ptr, ring_locs_ptr, @@ -84,6 +86,9 @@ def _qsa_graph_row_metadata_kernel( max_pages, RATIO: tl.constexpr, FULL_PAGE: tl.constexpr, # full-KV tokens per page + COMPRESSED_PAGE: tl.constexpr, + WORLD_SIZE: tl.constexpr, + RANK: tl.constexpr, PAGE_BLOCK: tl.constexpr, seq_lens_ptr, req_pool_ptr, @@ -114,18 +119,22 @@ def _qsa_graph_row_metadata_kernel( seq_len = tl.load(row_seq_lens_ptr + row).to(tl.int32) req = tl.load(row_req_pool_ptr + row).to(tl.int64) token_row = req * req_to_token_row_stride + compressed = seq_len // RATIO if tl.program_id(1) == 0: current = tl.maximum(seq_len - 1, 0) last_loc = tl.load(req_to_token_ptr + token_row + current).to(tl.int32) - compressed = seq_len // RATIO - tl.store(compressed_lens_ptr + row, compressed) - - # The page-aligned allocator keeps each compression group inside one page, - # so last_loc // RATIO is the group's compressed slot; slot 0 is the padding slot. + # The page-aligned allocator keeps each compression group inside one page. boundary = (seq_len > 0) & (seq_len % RATIO == 0) - write_loc = tl.where(boundary, last_loc // RATIO, 0) - tl.store(write_locs_ptr + row, write_loc) + global_write = last_loc // RATIO + global_write_page = global_write // COMPRESSED_PAGE + write_owner = boundary & (global_write_page % WORLD_SIZE == RANK) + local_write = ( + global_write_page // WORLD_SIZE * COMPRESSED_PAGE + + global_write % COMPRESSED_PAGE + ) + tl.store(write_locs_ptr + row, tl.where(write_owner, local_write, 0)) + tl.store(write_owner_ptr + row, write_owner) tl.store(logical_positions_ptr + row, current) tl.store(state_slots_ptr + row, req * RATIO + (current % RATIO).to(tl.int64)) @@ -139,12 +148,46 @@ def _qsa_graph_row_metadata_kernel( # page-aligned req_to_token row; the scoring kernels turn them into # compressed slots as page_id * (FULL_PAGE // RATIO) + block_in_page. table_row = page_table_ptr + row.to(tl.int64) * max_pages + positions_row = block_positions_ptr + row.to(tl.int64) * max_pages * COMPRESSED_PAGE offs = tl.arange(0, PAGE_BLOCK) row_width_pages = req_to_token_row_stride // FULL_PAGE - idx = tl.program_id(1) * PAGE_BLOCK + offs - valid = idx < tl.minimum(max_pages, row_width_pages) - loc = tl.load(req_to_token_ptr + token_row + idx * FULL_PAGE, mask=valid, other=0) - tl.store(table_row + idx, tl.maximum(loc // FULL_PAGE, 0), mask=valid) + visible_pages = tl.cdiv(compressed, COMPRESSED_PAGE) + local_page_count = 0 + local_block_count = 0 + for p0 in range(0, visible_pages, PAGE_BLOCK): + idx = p0 + offs + valid = idx < tl.minimum(visible_pages, tl.minimum(max_pages, row_width_pages)) + loc = tl.load( + req_to_token_ptr + token_row + idx * FULL_PAGE, mask=valid, other=0 + ) + global_page = tl.maximum(loc // FULL_PAGE, 0) + owner = valid & (global_page % WORLD_SIZE == RANK) + owner_i = owner.to(tl.int32) + local_indices = local_page_count + tl.cumsum(owner_i, axis=0) - 1 + tl.store( + table_row + local_indices, + global_page // WORLD_SIZE, + mask=owner, + ) + blocks = tl.maximum( + 0, + tl.minimum( + COMPRESSED_PAGE, + compressed - idx * COMPRESSED_PAGE, + ), + ) + for k in tl.static_range(COMPRESSED_PAGE): + block_valid = owner & (k < blocks) + tl.store( + positions_row + local_indices * COMPRESSED_PAGE + k, + tl.where(block_valid, idx * COMPRESSED_PAGE + k, -1), + mask=owner, + ) + local_page_count += tl.sum(owner_i, axis=0) + local_block_count += tl.sum(tl.where(owner, blocks, 0), axis=0) + # Consumers are bounded by local_block_count, so graph replays may leave the + # unreachable tail intact when a sequence shrinks. + tl.store(compressed_lens_ptr + row, local_block_count) def supports_graph_metadata_kernels(pool, device) -> bool: @@ -176,6 +219,7 @@ def launch_graph_metadata( row_seq_lens = metadata.sequence_lengths row_req_pool = metadata.row_req_pool_indices row_prefix_lens = indexer.graph_prefix_lengths + compressed_sharding = getattr(pool, "qsa_compressed_sharding", None) if mode == 2: _qsa_graph_layout_kernel[(bs + 1,)]( @@ -194,12 +238,14 @@ def launch_graph_metadata( num_padding, num_warps=1, ) - _qsa_graph_row_metadata_kernel[(num_rows, triton.cdiv(max_pages, 128))]( + _qsa_graph_row_metadata_kernel[(num_rows, 1)]( row_seq_lens, row_req_pool, indexer.graph_compressed_lengths, indexer.graph_write_locs, + indexer.write_owner_mask, indexer.graph_compressed_page_table, + indexer.decode_block_positions, indexer.decode_logical_positions, indexer.pending_ring_slots, indexer.graph_ring_group_locs, @@ -208,6 +254,11 @@ def launch_graph_metadata( max_pages, RATIO=indexer.compress_ratio, FULL_PAGE=pool.qsa_compressed_page_size * indexer.compress_ratio, + COMPRESSED_PAGE=pool.qsa_compressed_page_size, + WORLD_SIZE=( + compressed_sharding.world_size if compressed_sharding is not None else 1 + ), + RANK=(compressed_sharding.rank if compressed_sharding is not None else 0), PAGE_BLOCK=128, seq_lens_ptr=seq_lens, req_pool_ptr=req_pool_indices, @@ -233,6 +284,9 @@ def _qsa_draft_graph_metadata_kernel( MAX_PAGES: tl.constexpr, RATIO: tl.constexpr, FULL_PAGE: tl.constexpr, + COMPRESSED_PAGE: tl.constexpr, + WORLD_SIZE: tl.constexpr, + RANK: tl.constexpr, ): for step in tl.static_range(len(buffers)): if tl.program_id(2) == step: @@ -246,15 +300,20 @@ def _qsa_draft_graph_metadata_kernel( buf[5], buf[6], buf[7], + buf[8], + buf[9], req_to_token, row_stride, MAX_PAGES, RATIO, FULL_PAGE, + COMPRESSED_PAGE, + WORLD_SIZE, + RANK, 128, seq_lens, req_pool, - buf[8], + buf[10], bs, num_padding, extend_len=0, @@ -264,13 +323,16 @@ def _qsa_draft_graph_metadata_kernel( def prepare_draft_graph_metadata(metadata, req_to_token, pool): + compressed_sharding = getattr(pool, "qsa_compressed_sharding", None) buffers = tuple( ( m.sequence_lengths, m.row_req_pool_indices, m.indexer_metadata.graph_compressed_lengths, m.indexer_metadata.graph_write_locs, + m.indexer_metadata.write_owner_mask, m.indexer_metadata.graph_compressed_page_table, + m.indexer_metadata.decode_block_positions, m.indexer_metadata.decode_logical_positions, m.indexer_metadata.pending_ring_slots, m.indexer_metadata.graph_ring_group_locs, @@ -285,14 +347,27 @@ def prepare_draft_graph_metadata(metadata, req_to_token, pool): metadata[0].indexer_metadata.graph_compressed_page_table.shape[1], pool.qsa_compress_ratio, pool.qsa_compressed_page_size * pool.qsa_compress_ratio, + pool.qsa_compressed_page_size, + (compressed_sharding.world_size if compressed_sharding is not None else 1), + (compressed_sharding.rank if compressed_sharding is not None else 0), ) def launch_draft_graph_metadata(args, seq_lens, req_pool_indices, bs, num_padding): - buffers, req_to_token, row_stride, max_pages, ratio, full_page = args + ( + buffers, + req_to_token, + row_stride, + max_pages, + ratio, + full_page, + compressed_page, + world_size, + rank, + ) = args if bs == 0: return - _qsa_draft_graph_metadata_kernel[(bs, triton.cdiv(max_pages, 128), len(buffers))]( + _qsa_draft_graph_metadata_kernel[(bs, 1, len(buffers))]( buffers, seq_lens, req_pool_indices, @@ -303,5 +378,8 @@ def launch_draft_graph_metadata(args, seq_lens, req_pool_indices, bs, num_paddin max_pages, ratio, full_page, + compressed_page, + world_size, + rank, num_warps=1, ) diff --git a/python/sglang/srt/layers/attention/qsa/metadata.py b/python/sglang/srt/layers/attention/qsa/metadata.py index 2efd0998b781..646fbe321269 100644 --- a/python/sglang/srt/layers/attention/qsa/metadata.py +++ b/python/sglang/srt/layers/attention/qsa/metadata.py @@ -15,6 +15,10 @@ gather_qsa_compressed_keys, qsa_fast_topk, ) +from sglang.srt.mem_cache.qsa_kv_pool import ( + QSACompressedBlockSharding, + assert_qsa_indices_in_bounds, +) def build_qsa_row_ranges( @@ -46,6 +50,56 @@ def build_qsa_row_ranges( return row_starts, row_ends, compressed_cu_seqlens +def count_visible_local_blocks( + logical_block_positions: torch.Tensor, + local_lengths: torch.Tensor, + sequence_ids: torch.Tensor, + visible_blocks: torch.Tensor, + logical_position_capacity: int, +) -> torch.Tensor: + """Count visible owner-local blocks without a rows-by-prefix workspace.""" + + if logical_block_positions.ndim != 2: + raise ValueError("logical_block_positions must be rank 2") + if local_lengths.ndim != 1: + raise ValueError("local_lengths must be rank 1") + if logical_block_positions.shape[0] != local_lengths.numel(): + raise ValueError("logical block rows must match local lengths") + if sequence_ids.shape != visible_blocks.shape: + raise ValueError("sequence_ids and visible_blocks must have matching shapes") + if logical_position_capacity < 0: + raise ValueError("logical_position_capacity must be non-negative") + row_width = logical_block_positions.shape[1] + if row_width == 0: + return torch.zeros_like(visible_blocks, dtype=torch.int32) + + # Offset each sequence into a disjoint sorted key range. Invalid tail + # entries sort after all real positions in their row, so one searchsorted + # handles every query row without host loops or scalar readbacks. + stride = logical_position_capacity + 1 + metadata_rows = torch.arange( + logical_block_positions.shape[0], + device=logical_block_positions.device, + dtype=torch.int64, + ) + row_bases = metadata_rows * stride + sorted_keys = torch.where( + logical_block_positions >= 0, + logical_block_positions + row_bases.unsqueeze(1), + row_bases.unsqueeze(1) + logical_position_capacity, + ).flatten() + query_sequence_ids = sequence_ids.to( + device=logical_block_positions.device, dtype=torch.int64 + ) + query_keys = query_sequence_ids * stride + visible_blocks.to( + device=logical_block_positions.device, dtype=torch.int64 + ).clamp(min=0, max=logical_position_capacity) + absolute_counts = torch.searchsorted(sorted_keys, query_keys, right=False) + counts = absolute_counts - query_sequence_ids * row_width + row_lengths = local_lengths.index_select(0, query_sequence_ids) + return torch.minimum(counts, row_lengths).to(torch.int32) + + class QSAIndexerMetadata(msgspec.Struct, frozen=True): """All per-forward metadata consumed specifically by ``QSAIndexer``. @@ -78,6 +132,8 @@ class QSAIndexerMetadata(msgspec.Struct, frozen=True): # The first member's token row in this forward's packed tensors is extend only, # where group-aligned chunks keep every member in-chunk; None on paged forwards. write_locs: Optional[torch.Tensor] = None + write_owner_mask: Optional[torch.Tensor] = None + local_write_locs: Optional[torch.Tensor] = None compress_group_positions: Optional[torch.Tensor] = None compress_sequence_ids: Optional[torch.Tensor] = None compress_member_rows: Optional[torch.Tensor] = None @@ -86,8 +142,13 @@ class QSAIndexerMetadata(msgspec.Struct, frozen=True): graph_compressed_page_table: Optional[torch.Tensor] = None graph_compressed_lengths: Optional[torch.Tensor] = None graph_prefix_lengths: Optional[torch.Tensor] = None + prefill_page_table: Optional[torch.Tensor] = None + prefill_lengths: Optional[torch.Tensor] = None + prefill_block_positions: Optional[torch.Tensor] = None + prefill_cache_indices: Optional[torch.Tensor] = None decode_page_table: Optional[torch.Tensor] = None decode_lengths: Optional[torch.Tensor] = None + decode_block_positions: Optional[torch.Tensor] = None decode_logical_positions: Optional[torch.Tensor] = None pending_ring_slots: Optional[torch.Tensor] = None compress_group_ring_locs: Optional[torch.Tensor] = None @@ -158,10 +219,21 @@ def get_prefill_mqa_inputs( ratio = self.compress_ratio compressed_buffer = pool.get_qsa_compressed_k_buffer(layer_id) sequence_lengths = self.sequence_lengths.to(torch.int32) - sequence_lengths_list = self.prefill_sequence_lengths_cpu - if sequence_lengths_list is None: - sequence_lengths_list = sequence_lengths.tolist() - complete_blocks = [int(length) // ratio for length in sequence_lengths_list] + if self.prefill_page_table is not None: + if ( + self.prefill_lengths is None + or self.prefill_block_positions is None + or self.prefill_cache_indices is None + ): + raise RuntimeError("QSA local prefill page metadata is incomplete") + compressed_keys = compressed_buffer.index_select( + 0, self.prefill_cache_indices + ) + else: + sequence_lengths_list = self.prefill_sequence_lengths_cpu + if sequence_lengths_list is None: + sequence_lengths_list = sequence_lengths.tolist() + complete_blocks = [int(length) // ratio for length in sequence_lengths_list] if query_sequence_ids is None: query_sequence_ids = self.token_to_batch_idx num_valid_tokens = query_sequence_ids.numel() @@ -171,20 +243,41 @@ def get_prefill_mqa_inputs( f"positions={positions.numel()}, mapping={num_valid_tokens}" ) positions = positions[:num_valid_tokens] - row_starts, row_ends, compressed_cu_seqlens = build_qsa_row_ranges( - sequence_lengths, - positions.to(sequence_lengths.device), - query_sequence_ids.to(sequence_lengths.device), - self.compress_ratio, - ) - compressed_keys = gather_qsa_compressed_keys( - compressed_buffer, - self.token_slot_table, - compressed_cu_seqlens, - sum(complete_blocks), - max(complete_blocks, default=0), - ratio, - ) + if self.prefill_lengths is None: + row_starts, row_ends, compressed_cu_seqlens = build_qsa_row_ranges( + sequence_lengths, + positions.to(sequence_lengths.device), + query_sequence_ids.to(sequence_lengths.device), + self.compress_ratio, + ) + compressed_keys = gather_qsa_compressed_keys( + compressed_buffer, + self.token_slot_table, + compressed_cu_seqlens, + sum(complete_blocks), + max(complete_blocks, default=0), + ratio, + ) + else: + cumulative = torch.nn.functional.pad( + self.prefill_lengths.to(torch.int32).cumsum(0), (1, 0) + ).to(torch.int32) + sequence_ids = query_sequence_ids.long() + row_starts = cumulative.index_select(0, sequence_ids) + visible_blocks = torch.div( + positions.to(sequence_lengths.device, dtype=torch.int64) + 1, + ratio, + rounding_mode="floor", + ) + visible_local = count_visible_local_blocks( + self.prefill_block_positions, + self.prefill_lengths, + sequence_ids, + visible_blocks, + logical_position_capacity=(self.token_slot_table.shape[1] + ratio - 1) + // ratio, + ) + row_ends = row_starts + visible_local return compressed_keys, row_starts, row_ends, sequence_lengths def get_decode_mqa_inputs( @@ -295,6 +388,24 @@ def build_rope_position_matrix( return rope_positions[:, :num_tokens].long().transpose(0, 1).contiguous() +def build_qsa_capture_token_slot_table( + *, + batch_size: int, + max_length: int, + device: torch.device, +) -> torch.Tensor: + """Build cyclic physical pages for synthetic prefill graph requests.""" + if batch_size < 0: + raise ValueError("batch_size must be non-negative") + if max_length < 0: + raise ValueError("max_length must be non-negative") + return ( + torch.arange(max_length, dtype=torch.int32, device=device) + .unsqueeze(0) + .expand(batch_size, -1) + ) + + def compressed_decode_view( *, compressed_page_size: int, @@ -324,11 +435,177 @@ def compressed_decode_view( return compressed_page_table, compressed_lengths +def localize_compressed_page_table( + *, + global_page_table: torch.Tensor, + compressed_lengths: torch.Tensor, + sharding: QSACompressedBlockSharding, +) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Localize owned physical pages and preserve each logit's logical position. + + Page-table rows may contain non-monotonic physical page ids. Ownership and + cache addressing follow those physical ids, while ``logical_positions`` + records the original sequence-local compressed-block position represented + by each compacted logit. + """ + + if global_page_table.ndim != 2: + raise ValueError( + "global_page_table must be rank 2, got " + f"shape={tuple(global_page_table.shape)}" + ) + if compressed_lengths.ndim != 1: + raise ValueError( + "compressed_lengths must be rank 1, got " + f"shape={tuple(compressed_lengths.shape)}" + ) + if global_page_table.shape[0] != compressed_lengths.numel(): + raise ValueError( + "page-table rows must match compressed lengths: " + f"rows={global_page_table.shape[0]}, " + f"lengths={compressed_lengths.numel()}" + ) + if global_page_table.device != compressed_lengths.device: + raise ValueError("global_page_table and compressed_lengths must share a device") + non_negative = torch.all(compressed_lengths >= 0) + if compressed_lengths.is_cuda: + torch._assert_async(non_negative, "compressed_lengths must be non-negative") + elif not bool(non_negative.item()): + raise ValueError("compressed_lengths must be non-negative") + + page_size = sharding.compressed_page_size + row_capacity = global_page_table.shape[1] * page_size + within_capacity = torch.all(compressed_lengths <= row_capacity) + if compressed_lengths.is_cuda: + torch._assert_async( + within_capacity, + f"compressed length exceeds page-table capacity: capacity={row_capacity}", + ) + elif not bool(within_capacity.item()): + raise ValueError( + f"compressed length exceeds page-table capacity: capacity={row_capacity}" + ) + + logical_page_starts = ( + torch.arange( + global_page_table.shape[1], + dtype=torch.int64, + device=global_page_table.device, + ) + * page_size + ) + visible_pages = logical_page_starts.unsqueeze( + 0 + ) < compressed_lengths.long().unsqueeze(1) + visible_physical_pages = global_page_table[visible_pages] + if not global_page_table.is_cuda and torch.any( + (visible_physical_pages < 0) | (visible_physical_pages >= sharding.global_pages) + ): + raise ValueError("visible compressed physical page is out of range") + assert_qsa_indices_in_bounds( + global_page_table, + sharding.global_pages, + valid_mask=visible_pages, + label="compressed KV page table", + ) + + local_page_ids = sharding.global_to_local_pages(global_page_table) + owned_pages = visible_pages & (local_page_ids >= 0) + assert_qsa_indices_in_bounds( + local_page_ids, + sharding.local_pages, + valid_mask=owned_pages, + label="compressed KV local page table", + ) + local_page_counts = owned_pages.sum(dim=1, dtype=torch.int32) + logical_block_positions = ( + logical_page_starts[:, None] + + torch.arange( + page_size, + dtype=torch.int64, + device=global_page_table.device, + )[None, :] + ) + visible_blocks = logical_block_positions.unsqueeze( + 0 + ) < compressed_lengths.long().view(-1, 1, 1) + local_lengths = (visible_blocks & owned_pages.unsqueeze(-1)).sum( + dim=(1, 2), dtype=torch.int32 + ) + + output_pages = ( + global_page_table.shape[1] + sharding.world_size - 1 + ) // sharding.world_size + within_local_capacity = torch.all(local_page_counts <= output_pages) + if global_page_table.is_cuda: + torch._assert_async( + within_local_capacity, + "compressed page ownership must follow the allocator's cyclic layout", + ) + elif not bool(within_local_capacity.item()): + raise ValueError("compressed page ownership is not cyclic") + + compact_page_mask = torch.arange( + output_pages, device=global_page_table.device + ).unsqueeze(0) < local_page_counts.unsqueeze(1) + local_page_table = torch.full( + compact_page_mask.shape, + -1, + dtype=global_page_table.dtype, + device=global_page_table.device, + ) + compacted_columns = torch.full( + compact_page_mask.shape, + -1, + dtype=torch.int64, + device=global_page_table.device, + ) + if output_pages: + destinations = owned_pages.cumsum(dim=1).sub(1).clamp_min(0).long() + local_page_table.scatter_reduce_( + 1, + destinations, + torch.where(owned_pages, local_page_ids, -1).to(global_page_table.dtype), + reduce="amax", + include_self=True, + ) + source_columns = torch.arange( + global_page_table.shape[1], + dtype=torch.int64, + device=global_page_table.device, + ).expand_as(global_page_table) + compacted_columns.scatter_reduce_( + 1, + destinations, + torch.where(owned_pages, source_columns, -1), + reduce="amax", + include_self=True, + ) + + compact_logical_positions = compacted_columns.unsqueeze( + -1 + ) * page_size + torch.arange( + page_size, + dtype=torch.int64, + device=global_page_table.device, + ).view(1, 1, -1) + compact_block_mask = compact_page_mask.unsqueeze(-1) & ( + compact_logical_positions < compressed_lengths.long().view(-1, 1, 1) + ) + logical_positions = torch.where( + compact_block_mask, compact_logical_positions, -1 + ).flatten(1) + return local_page_table, local_lengths, logical_positions + + __all__ = [ "QSAIndexerMetadata", + "build_qsa_capture_token_slot_table", "build_qsa_row_ranges", "build_pending_ring_slots", "build_group_ring_slots", "build_rope_position_matrix", "compressed_decode_view", + "count_visible_local_blocks", + "localize_compressed_page_table", ] diff --git a/python/sglang/srt/layers/attention/qsa/owner_sparse_attn.py b/python/sglang/srt/layers/attention/qsa/owner_sparse_attn.py new file mode 100644 index 000000000000..2ecd5b208be5 --- /dev/null +++ b/python/sglang/srt/layers/attention/qsa/owner_sparse_attn.py @@ -0,0 +1,251 @@ +from __future__ import annotations + +import torch +import triton +import triton.language as tl + +from sglang.srt.distributed.utils import all_gather_single +from sglang.srt.mem_cache.qsa_kv_pool import ( + QSARawKVSharding, + assert_qsa_indices_in_bounds, +) + + +@triton.jit +def _owner_localize_token_slots_kernel( + global_slots_ptr, + compact_ptr, + lengths_ptr, + selected_width: tl.constexpr, + compact_width: tl.constexpr, + global_size: tl.constexpr, + world_size: tl.constexpr, + owner_rank: tl.constexpr, + block_width: tl.constexpr, +): + row = tl.program_id(0) + offsets = tl.arange(0, block_width) + input_mask = offsets < selected_width + global_ids = tl.load( + global_slots_ptr + row * selected_width + offsets, + mask=input_mask, + other=-1, + ).to(tl.int64) + owned = ( + input_mask + & (global_ids >= 0) + & (global_ids < global_size) + & (global_ids % world_size == owner_rank) + ) + destinations = tl.cumsum(owned.to(tl.int32), axis=0) - 1 + output_mask = offsets < compact_width + tl.store(compact_ptr + row * compact_width + offsets, -1, mask=output_mask) + tl.store( + compact_ptr + row * compact_width + destinations, + global_ids // world_size, + mask=owned & (destinations < compact_width), + ) + tl.store(lengths_ptr + row, tl.sum(owned.to(tl.int32), axis=0)) + + +def _owner_localize_token_slots_triton( + global_token_slots: torch.Tensor, + sharding: QSARawKVSharding, + *, + fixed_width: int, + compact: torch.Tensor | None = None, + lengths: torch.Tensor | None = None, +) -> tuple[torch.Tensor, torch.Tensor]: + rows, selected_width = global_token_slots.shape + if compact is None: + compact = torch.empty( + (rows, fixed_width), + dtype=torch.int64, + device=global_token_slots.device, + ) + if lengths is None: + lengths = torch.empty( + (rows,), dtype=torch.int32, device=global_token_slots.device + ) + block_width = triton.next_power_of_2(max(selected_width, fixed_width)) + _owner_localize_token_slots_kernel[(rows,)]( + global_token_slots, + compact, + lengths, + selected_width=selected_width, + compact_width=fixed_width, + global_size=sharding.global_capacity, + world_size=sharding.world_size, + owner_rank=sharding.rank, + block_width=block_width, + num_warps=4 if block_width <= 128 else 8, + ) + if fixed_width < selected_width: + torch._assert_async( + torch.all(lengths <= fixed_width), + "owner token count exceeds interleaved fixed-width capacity", + ) + return compact, lengths + + +def select_kv_heads_for_query_shard( + k: torch.Tensor, + v: torch.Tensor, + *, + total_query_heads: int, + total_kv_heads: int, + tp_size: int, + tp_rank: int, +) -> tuple[torch.Tensor, torch.Tensor]: + if k.ndim != 3 or v.shape != k.shape: + raise ValueError("K/V must have matching [tokens, heads, dim] shapes") + if ( + min(total_query_heads, total_kv_heads, tp_size) <= 0 + or total_query_heads % total_kv_heads + or total_query_heads % tp_size + or not 0 <= tp_rank < tp_size + ): + raise ValueError("invalid query/KV tensor-parallel head geometry") + query_heads_per_rank = total_query_heads // tp_size + queries_per_kv = total_query_heads // total_kv_heads + query_start = tp_rank * query_heads_per_rank + query_stop = query_start + query_heads_per_rank + first_kv = query_start // queries_per_kv + last_kv = (query_stop - 1) // queries_per_kv + if first_kv != last_kv: + raise ValueError( + "query shard crosses KV-head groups; QSA requires group-aligned TP" + ) + if k.shape[1] == 1: + return k, v + if k.shape[1] != total_kv_heads: + raise ValueError( + "K/V head layout must be either owner-local or globally replicated: " + f"got {k.shape[1]}, expected 1 or {total_kv_heads}" + ) + return ( + k[:, first_kv : first_kv + 1].contiguous(), + v[:, first_kv : first_kv + 1].contiguous(), + ) + + +def owner_sparse_attention( + q: torch.Tensor, + local_k_cache: torch.Tensor, + local_v_cache: torch.Tensor, + global_token_slots: torch.Tensor, + sharding: QSARawKVSharding, + *, + group=None, + softmax_scale: float | None = None, + static_buffers=None, +) -> tuple[torch.Tensor, torch.Tensor]: + if static_buffers is not None: + static_buffers = static_buffers.for_batch(q.shape[0]) + if q.shape[1] != static_buffers.num_heads // static_buffers.world_size: + raise ValueError("static Q transport head geometry mismatch") + gathered_q = q + if group is not None: + with torch.profiler.record_function("qsa.owner.q_gather"): + if static_buffers is None: + gathered_q = group.all_gather(q.contiguous(), dim=1).contiguous() + else: + rows = q.shape[0] + static_buffers.q_send.zero_() + static_buffers.q_send[:, :rows].copy_(q.permute(1, 0, 2)) + device_group = getattr(group, "device_group", None) + if device_group is None: + group.all_gather_into_tensor( + static_buffers.q_recv, static_buffers.q_send + ) + else: + work = all_gather_single( + static_buffers.q_recv, + static_buffers.q_send, + group=device_group, + async_op=True, + ) + work.block_current_stream() + with torch.profiler.record_function("qsa.owner.localize_slots"): + owner_width = global_token_slots.shape[1] + if static_buffers is not None and global_token_slots.is_cuda: + if owner_width > static_buffers.owner_topk: + raise ValueError("selected owner width exceeds static buffer capacity") + local_slots, _ = _owner_localize_token_slots_triton( + global_token_slots, + sharding, + fixed_width=owner_width, + compact=static_buffers.local_slots[ + : global_token_slots.shape[0], :owner_width + ], + lengths=static_buffers.owner_lengths[: global_token_slots.shape[0]], + ) + else: + local_slots, _ = _owner_localize_token_slots_triton( + global_token_slots, sharding, fixed_width=owner_width + ) + if static_buffers is not None and group is not None: + static_buffers.q_gathered[: q.shape[0]].copy_( + static_buffers.q_recv[:, : q.shape[0]].permute(1, 0, 2) + ) + gathered_q = static_buffers.q_gathered[: q.shape[0]] + with torch.profiler.record_function("qsa.owner.kernel"): + assert_qsa_indices_in_bounds( + local_slots, + local_k_cache.shape[0], + valid_mask=local_slots >= 0, + label="owner local KV read locations", + ) + from sglang.srt.layers.attention.qsa.sparse_attn import ( + sparse_gqa_packed_decode_with_lse_triton, + ) + + scale = gathered_q.shape[-1] ** -0.5 if softmax_scale is None else softmax_scale + output, lse = sparse_gqa_packed_decode_with_lse_triton( + gathered_q.contiguous(), + local_k_cache, + local_v_cache, + local_slots, + scale, + ) + if group is None: + return output, lse + with torch.profiler.record_function("qsa.owner.output_lse_merge"): + output_dtype = output.dtype + world_size = int(group.world_size) + rows, num_heads, head_dim = output.shape + if num_heads % world_size: + raise ValueError("owner partial heads must be divisible by group size") + local_heads = num_heads // world_size + if static_buffers is None: + packed = torch.empty( + (world_size, rows, local_heads, head_dim + 1), + dtype=torch.float32, + device=output.device, + ) + recv = torch.empty_like(packed) + else: + packed = static_buffers.owner_a2a_send + recv = static_buffers.owner_a2a_recv + packed[..., :-1].zero_() + packed[..., -1].fill_(float("-inf")) + packed[:, :rows, ..., :-1].copy_( + output.float() + .view(rows, world_size, local_heads, head_dim) + .permute(1, 0, 2, 3) + ) + packed[:, :rows, ..., -1].copy_( + lse.view(rows, world_size, local_heads).permute(1, 0, 2) + ) + group.all_to_all_single(recv.view(-1), packed.view(-1)) + gathered = recv[:, :rows] + from sglang.kernels.ops.attention.dcp_kernels import dcp_lse_combine_triton + + output, lse = dcp_lse_combine_triton( + gathered[..., :-1], + gathered[..., -1], + is_lse_base_on_e=True, + return_lse=True, + ) + output = output.to(output_dtype) + return output, lse diff --git a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py index a550ff7d7001..d6252da54dce 100644 --- a/python/sglang/srt/layers/attention/qsa/qsa_indexer.py +++ b/python/sglang/srt/layers/attention/qsa/qsa_indexer.py @@ -2,11 +2,20 @@ from __future__ import annotations +import os from typing import Optional, Tuple import torch from sglang.srt.batch_invariant_ops import is_batch_invariant_mode_enabled +from sglang.srt.layers.attention.qsa.cache_sharding import ( + QSACacheShardingRuntime, + assert_qsa_cache_sharding_runtime_match, + get_qsa_cache_sharding_runtime, +) +from sglang.srt.layers.attention.qsa.distributed_topk import ( + gather_and_merge_qsa_topk_candidates, +) from sglang.srt.layers.attention.qsa.kernel import ( average_pool_qsa_keys, expand_qsa_block_indices, @@ -34,6 +43,7 @@ from sglang.srt.layers.linear import ReplicatedLinear from sglang.srt.layers.rotary_embedding.utils import apply_rotary_emb from sglang.srt.layers.utils.multi_platform import MultiPlatformOp +from sglang.srt.mem_cache.qsa_kv_pool import assert_qsa_indices_in_bounds from sglang.srt.model_executor.runner import get_is_capture_mode # Cap on the fp32 [query_rows, compressed_keys] prefill logits workspace; @@ -41,9 +51,23 @@ _QSA_PREFILL_LOGITS_BUDGET_BYTES = 128 * 1024 * 1024 -def _qsa_prefill_row_chunk_size(rows: int, keys: int, heads: int) -> int: +def _qsa_prefill_row_chunk_size( + rows: int, + keys: int, + heads: int, + *, + collective_keys: int | None = None, + collective_world_size: int = 1, +) -> int: if rows <= 0 or keys <= 0: return max(rows, 1) + if collective_keys is not None: + if collective_world_size <= 0: + raise ValueError("collective_world_size must be positive") + # Every rank must issue the same sequence of candidate collectives. + # Rank-local sharded cache lengths can differ, so derive the workspace + # schedule from the shared global key count instead. + keys = -(collective_keys // -collective_world_size) block_q = max(1, 128 // heads) bytes_per_row = keys * torch.float32.itemsize max_padded_rows = max(block_q, _QSA_PREFILL_LOGITS_BUDGET_BYTES // bytes_per_row) @@ -61,9 +85,22 @@ def __init__( quant_config=None, prefix: str = "", rotary_emb=None, + cache_sharding_runtime: QSACacheShardingRuntime | None = None, ) -> None: self._validate_config(config) super().__init__() + resolved_runtime = get_qsa_cache_sharding_runtime() + if cache_sharding_runtime is None: + cache_sharding_runtime = resolved_runtime + assert_qsa_cache_sharding_runtime_match( + resolved_runtime, + cache_sharding_runtime, + component="QSA indexer", + ) + expected_group = ( + cache_sharding_runtime.group if cache_sharding_runtime.enabled else None + ) + self.cache_sharding_runtime = cache_sharding_runtime self.layer_id = int(layer_id) self.index_n_heads = int(config.indexer_n_heads) self.index_kv_heads = int(config.indexer_kv_heads) @@ -93,6 +130,7 @@ def __init__( self.index_head_dim, eps=getattr(config, "rms_norm_eps", 1e-6) ) self._rope_axis_map_cache = None + self.distributed_topk_group = expected_group @staticmethod def _validate_config(config) -> None: @@ -195,13 +233,20 @@ def project_qk( int(positions.max().item()) ) key_state_buffer = pool.get_qsa_key_state_buffer(self.layer_id) + fused_state_slots = cache_loc[: qk.shape[0]].long() + assert_qsa_indices_in_bounds( + fused_state_slots, + key_state_buffer.shape[0], + valid_mask=torch.ones_like(fused_state_slots, dtype=torch.bool), + label="QSA pending-state write locations", + ) q = qsa_index_q_norm_rope_store( qk, positions.long(), self.rotary_emb.cos_sin_cache, self._rope_axis_map(qk.device), self.q_layernorm.weight.data, - cache_loc[: qk.shape[0]].long(), + fused_state_slots, key_state_buffer.view(key_state_buffer.shape[0], -1), pool.qsa_rope_position_buffer, self.index_n_heads, @@ -240,7 +285,8 @@ def _fused_compress_store( self, pool, group_locs: torch.Tensor, - write_locs: torch.Tensor, + owner_mask: torch.Tensor, + local_write_locs: torch.Tensor | None = None, source_keys: torch.Tensor | None = None, source_rope: torch.Tensor | None = None, ) -> None: @@ -254,7 +300,23 @@ def _fused_compress_store( source_keys = pool.get_qsa_key_state_buffer(self.layer_id) if source_rope is None: source_rope = pool.qsa_rope_position_buffer + if local_write_locs is None: + local_write_locs = owner_mask + owner_mask = torch.ones_like(local_write_locs, dtype=torch.bool) compressed_buffer = pool.get_qsa_compressed_k_buffer(self.layer_id) + owner_rows = owner_mask.to(torch.bool) + assert_qsa_indices_in_bounds( + group_locs, + source_keys.shape[0], + valid_mask=owner_rows.unsqueeze(1).expand_as(group_locs), + label="QSA compressed source read locations", + ) + assert_qsa_indices_in_bounds( + local_write_locs, + compressed_buffer.shape[0], + valid_mask=owner_rows, + label="QSA compressed destination write locations", + ) qsa_index_k_compress_store( source_keys.reshape(source_keys.shape[0], -1) .contiguous() @@ -264,7 +326,8 @@ def _fused_compress_store( self.rotary_emb.cos_sin_cache, self._rope_axis_map(source_keys.device), self.k_layernorm.weight.data, - write_locs.to(torch.int32), + owner_mask.to(torch.int32), + local_write_locs.to(torch.int32), compressed_buffer.view(compressed_buffer.shape[0], -1), self.compress_ratio, self.rotary_emb.rotary_dim, @@ -331,7 +394,14 @@ def update_key_state_and_compress( if metadata.write_locs.numel() == 0: return group_end_positions = metadata.compress_group_positions.long() - compressed_locs = metadata.write_locs + owner_mask = metadata.write_owner_mask + compressed_locs = metadata.local_write_locs + if owner_mask is None or compressed_locs is None: + raise RuntimeError( + "QSA eager compression metadata requires owner mask and local " + "write locations" + ) + owner_mask = owner_mask.bool() if is_extend: # Extend chunks are group-aligned; each planned group lies in this forward, # so read its members from the packed chunk tensors. @@ -362,30 +432,44 @@ def update_key_state_and_compress( self._fused_compress_store( pool, group_locs, + owner_mask, compressed_locs, source_keys=source_keys, source_rope=source_rope, ) return - key_groups = source_keys[group_locs] + owned_rows = torch.nonzero(owner_mask, as_tuple=False).flatten() + if owned_rows.numel() == 0: + return + owned_group_locs = group_locs.index_select(0, owned_rows) + key_groups = source_keys[owned_group_locs] pooled = average_pool_qsa_keys(key_groups) compressed_rope_positions = self._rope_from_matrix( - source_rope[group_locs[:, 0]] + source_rope[owned_group_locs[:, 0]] ) normalized = self.normalize_compressed_keys(pooled, compressed_rope_positions) - pool.set_qsa_compressed_k_buffer(self.layer_id, compressed_locs, normalized) + buffer = pool.get_qsa_compressed_k_buffer(self.layer_id) + buffer[compressed_locs.index_select(0, owned_rows).long()] = normalized.to( + buffer.dtype + ) def _compress_decode_cuda_graph(self, metadata) -> None: - """Fixed-shape graph-replay compression; non-boundary rows write slot 0.""" + """Fixed-shape compression with owner-masked writes to local slots.""" if metadata.graph_write_locs is None or metadata.graph_ring_group_locs is None: raise RuntimeError("QSA CUDA graph compression metadata is incomplete") pool = metadata.token_to_kv_pool group_locs = metadata.graph_ring_group_locs + owner_mask = ( + metadata.write_owner_mask + if metadata.write_owner_mask is not None + else metadata.graph_write_locs != 0 + ) if self._use_fused_compress(pool): self._fused_compress_store( pool, group_locs, + owner_mask, metadata.graph_write_locs, ) return @@ -397,9 +481,20 @@ def _compress_decode_cuda_graph(self, metadata) -> None: compressed = self.normalize_compressed_keys( compressed, compressed_rope_positions ) - pool.set_qsa_compressed_k_buffer( - self.layer_id, metadata.graph_write_locs, compressed.contiguous() - ) + # Graph metadata already contains local slots. The global-address pool + # setter would localize twice and dynamically compact the owner rows. + buffer = pool.get_qsa_compressed_k_buffer(self.layer_id) + if compressed.is_cuda: + from sglang.kernels.ops.attention.qsa_indexer import ( + qsa_store_compressed_keys, + ) + + qsa_store_compressed_keys( + compressed, buffer, metadata.graph_write_locs, owner_mask + ) + else: + local_locs = metadata.graph_write_locs[owner_mask].long() + buffer[local_locs] = compressed[owner_mask].to(buffer.dtype) def _rope_from_matrix(self, positions: torch.Tensor) -> torch.Tensor: """[n, 3] slot coordinates -> the layout apply_rope expects.""" @@ -448,6 +543,95 @@ def apply_rope( ) return torch.cat([rotated, tensor[..., rotary_dim:]], dim=-1) + def _maybe_merge_distributed_topk( + self, + logits: torch.Tensor | None, + local_indices: torch.Tensor, + row_starts: torch.Tensor, + logical_block_positions: torch.Tensor | None = None, + static_buffers=None, + ) -> torch.Tensor: + """Merge rank-local candidates using their explicit logical block ids.""" + + group = getattr(self, "distributed_topk_group", None) + if group is None: + return local_indices + if logical_block_positions is None: + raise ValueError("distributed QSA top-k requires logical block positions") + if ( + static_buffers is not None + and getattr(static_buffers, "candidate_symm_state", None) is not None + and logits is not None + and logits.shape[1] > 0 + and os.environ.get("SGLANG_QSA_FUSED_CANDIDATE_PUBLISH", "1") != "0" + ): + from sglang.srt.layers.attention.qsa import distributed_topk + + _, merged_indices = ( + distributed_topk.fused_publish_and_merge_qsa_topk_candidates( + logits, + local_indices, + row_starts, + logical_block_positions, + group=group, + topk=self.block_topk, + static_buffers=static_buffers, + ) + ) + return merged_indices + valid = local_indices >= 0 + if logical_block_positions.shape[1] == 0: + valid = torch.zeros_like(valid) + mapped_logical_positions = torch.full_like( + local_indices, -1, dtype=torch.int64 + ) + else: + within_local_row = local_indices < logical_block_positions.shape[1] + safe_local_indices = torch.where( + valid & within_local_row, + local_indices, + torch.zeros_like(local_indices), + ).long() + mapped_logical_positions = logical_block_positions.gather( + 1, safe_local_indices + ) + valid &= within_local_row & (mapped_logical_positions >= 0) + if logits is None or logits.shape[1] == 0: + local_scores = torch.full( + local_indices.shape, + float("-inf"), + dtype=torch.float32, + device=local_indices.device, + ) + else: + absolute_indices = local_indices.clamp_min( + 0 + ).long() + row_starts.long().unsqueeze(1) + within_logits = absolute_indices < logits.shape[1] + valid &= within_logits + safe_absolute_indices = torch.where( + valid, absolute_indices, torch.zeros_like(absolute_indices) + ) + local_scores = logits.gather(1, safe_absolute_indices).masked_fill( + ~valid, float("-inf") + ) + global_indices = torch.where( + valid, + mapped_logical_positions, + torch.full_like(mapped_logical_positions, -1), + ).to(torch.int32) + graph_kwargs = ( + {"static_buffers": static_buffers} if static_buffers is not None else {} + ) + _, merged_indices = gather_and_merge_qsa_topk_candidates( + local_scores, + global_indices, + group=group, + topk=self.block_topk, + **graph_kwargs, + ) + return merged_indices + def select_prefill_tokens( self, q: torch.Tensor, @@ -456,6 +640,11 @@ def select_prefill_tokens( row_ends: torch.Tensor, query_positions: torch.Tensor, sequence_lengths_for_rows: torch.Tensor, + logical_block_positions: torch.Tensor | None = None, + logical_block_sequence_ids: torch.Tensor | None = None, + static_buffers=None, + collective_keys: int | None = None, + collective_world_size: int = 1, max_row_keys: Optional[int] = None, ) -> torch.Tensor: """Top-k token selection for packed prefill rows. @@ -477,12 +666,18 @@ def select_prefill_tokens( return output relative = ( - max_row_keys is not None and max_row_keys > 0 and q.is_cuda and HAS_TILELANG + logical_block_positions is None + and max_row_keys is not None + and max_row_keys > 0 + and q.is_cuda + and HAS_TILELANG ) row_chunk_size = _qsa_prefill_row_chunk_size( rows, max_row_keys if relative else compressed_keys.shape[0], q.shape[1], + collective_keys=collective_keys, + collective_world_size=collective_world_size, ) for row_start in range(0, rows, row_chunk_size): row_end = min(row_start + row_chunk_size, rows) @@ -524,6 +719,23 @@ def select_prefill_tokens( row_ends[chunk_slice], topk=self.block_topk, ) + chunk_logical_block_positions = None + if logical_block_positions is not None: + chunk_logical_block_positions = ( + logical_block_positions[chunk_slice] + if logical_block_sequence_ids is None + else logical_block_positions.index_select( + 0, logical_block_sequence_ids[chunk_slice].long() + ) + ) + block_indices = QSAIndexer._maybe_merge_distributed_topk( + self, + logits, + block_indices, + row_starts[chunk_slice], + chunk_logical_block_positions, + static_buffers, + ) selected = expand_qsa_block_indices( block_indices, query_positions[chunk_slice], @@ -545,22 +757,29 @@ def select_decode_tokens( query_positions: torch.Tensor, sequence_lengths: torch.Tensor, defer_expansion: bool = False, + logical_block_positions: torch.Tensor | None = None, + static_buffers=None, ) -> torch.Tensor: - # fast_topk (k = 512) writes 0..length-1 for rows that fit k without reading - # their scores, and never reads past a row's length. + row_starts = torch.zeros_like(compressed_lengths, dtype=torch.int32) + # fast_topk (k = 512) writes 0..length-1 for rows that fit k without + # reading their scores, and never reads past a row's length. use_fast_topk = ( q.is_cuda and self.block_topk == 512 and not is_batch_invariant_mode_enabled() ) - logits = qsa_mqa_decode( - q, - compressed_cache, - compressed_page_table, - compressed_lengths, - max_model_len, - min_scored_len=self.block_topk if use_fast_topk else 0, + skip_local_scores = ( + use_fast_topk and getattr(self, "distributed_topk_group", None) is None ) + with torch.profiler.record_function("qsa.local_scan"): + logits = qsa_mqa_decode( + q, + compressed_cache, + compressed_page_table, + compressed_lengths, + max_model_len, + min_scored_len=self.block_topk if skip_local_scores else 0, + ) if use_fast_topk: # Decode rows start at zero, so compressed lengths double as row lengths; # skip the generic zero-fill + subtract. @@ -573,19 +792,27 @@ def select_decode_tokens( row_starts=None, ) else: - row_starts = torch.zeros_like(compressed_lengths, dtype=torch.int32) block_indices = qsa_fast_topk( logits, row_starts, compressed_lengths, topk=self.block_topk ) + block_indices = QSAIndexer._maybe_merge_distributed_topk( + self, + logits, + block_indices, + row_starts, + logical_block_positions, + static_buffers, + ) if defer_expansion: return block_indices - return expand_qsa_block_indices( + selected = expand_qsa_block_indices( block_indices, query_positions, sequence_lengths, compress_ratio=self.compress_ratio, token_topk=self.token_topk, ) + return selected def forward_cuda_cp( self, @@ -772,6 +999,15 @@ def _forward_impl( state_slots=state_slots, state_stored=state_stored, ) + static_buffers = ( + getattr( + indexer_metadata.token_to_kv_pool, + "qsa_cache_sharding_buffers", + None, + ) + if indexer_metadata.is_cuda_graph + else None + ) if forward_mode.is_decode() or is_target_verify or is_draft_extend: compressed_cache, page_table, compressed_lengths, max_model_len = ( indexer_metadata.get_decode_mqa_inputs(self.layer_id) @@ -785,12 +1021,39 @@ def _forward_impl( logical_positions, indexer_metadata.get_seqlens_int32(), defer_expansion=q.is_cuda and indexer_metadata.defer_block_expansion, + logical_block_positions=getattr( + indexer_metadata, "decode_block_positions", None + ), + static_buffers=static_buffers, ) compressed_keys, row_starts, row_ends, sequence_lengths = ( indexer_metadata.get_prefill_mqa_inputs(self.layer_id, logical_positions) ) query_sequence_ids = indexer_metadata.get_token_to_batch_idx() + prefill_kwargs = {} + block_positions = getattr(indexer_metadata, "prefill_block_positions", None) + if block_positions is not None: + prefill_kwargs["logical_block_positions"] = block_positions + prefill_kwargs["logical_block_sequence_ids"] = query_sequence_ids + group = getattr(self, "distributed_topk_group", None) + if group is not None: + host_lengths = indexer_metadata.prefill_sequence_lengths_cpu + if host_lengths is None: + # Hand-built metadata may not carry scheduler lengths. A + # shape-derived upper bound keeps every rank's collective + # schedule identical without synchronizing a CUDA scalar. + collective_keys = ( + indexer_metadata.sequence_lengths.numel() + * indexer_metadata.token_slot_table.shape[1] + // self.compress_ratio + ) + else: + collective_keys = sum( + length // self.compress_ratio for length in host_lengths + ) + prefill_kwargs["collective_keys"] = collective_keys + prefill_kwargs["collective_world_size"] = int(group.world_size) row_sequence_lengths = sequence_lengths.index_select( 0, query_sequence_ids.long() ) @@ -801,6 +1064,8 @@ def _forward_impl( row_ends, logical_positions, row_sequence_lengths, + **prefill_kwargs, + static_buffers=static_buffers, max_row_keys=indexer_metadata.prefill_max_row_keys( getattr(forward_batch, "seq_lens_cpu", None) ), diff --git a/python/sglang/srt/layers/attention/qsa/sparse_attn.py b/python/sglang/srt/layers/attention/qsa/sparse_attn.py index ec86b8f1304e..6a8e27bab4b3 100644 --- a/python/sglang/srt/layers/attention/qsa/sparse_attn.py +++ b/python/sglang/srt/layers/attention/qsa/sparse_attn.py @@ -174,6 +174,7 @@ def _sparse_gqa_chunk_prefill( k, v, out, + lse, indices, cu_q, cu_k, @@ -192,6 +193,8 @@ def _sparse_gqa_chunk_prefill( so_m: tl.constexpr, so_h: tl.constexpr, so_d: tl.constexpr, + sl_m: tl.constexpr, + sl_h: tl.constexpr, si_m: tl.constexpr, si_g: tl.constexpr, si_n: tl.constexpr, @@ -200,6 +203,7 @@ def _sparse_gqa_chunk_prefill( BLOCK_M: tl.constexpr, BLOCK_N: tl.constexpr, HEAD_DIM: tl.constexpr, + RETURN_LSE: tl.constexpr, SLOT_INDICES: tl.constexpr, ): """One query row x one KV head group over its ``topk`` selected tokens. @@ -272,7 +276,12 @@ def _sparse_gqa_chunk_prefill( ) normalizer = normalizer * alpha + tl.sum(probabilities, 1) max_value = next_max - output = accumulator / normalizer[:, None] + has_tokens = normalizer > 0 + output = tl.where( + has_tokens[:, None], + accumulator / tl.maximum(normalizer[:, None], 1.0), + 0.0, + ) tl.store( out + query * so_m @@ -281,6 +290,17 @@ def _sparse_gqa_chunk_prefill( output, mask=(offs_h < GROUP_SIZE)[:, None], ) + if RETURN_LSE: + row_lse = tl.where( + has_tokens, + (max_value + tl.log2(normalizer)) * 0.6931471805599453, + -float("inf"), + ) + tl.store( + lse + query * sl_m + (group * GROUP_SIZE + offs_h) * sl_h, + row_lse, + mask=offs_h < GROUP_SIZE, + ) @triton.jit(do_not_specialize=["TOTAL_K", "BS"]) @@ -400,6 +420,7 @@ def sparse_gqa_fwd_interface_triton_ck( k, v, out, + out, indices, cu_q, cu_k, @@ -418,6 +439,8 @@ def sparse_gqa_fwd_interface_triton_ck( out.stride(0), out.stride(1), out.stride(2), + 0, + 0, indices.stride(0), indices.stride(1) if indices.ndim == 3 else 0, indices.stride(2) if indices.ndim == 3 else indices.stride(1), @@ -426,6 +449,7 @@ def sparse_gqa_fwd_interface_triton_ck( BLOCK_M=block_m, BLOCK_N=block_n, HEAD_DIM=head_dim, + RETURN_LSE=False, SLOT_INDICES=False, num_warps=warps, num_stages=stages, @@ -453,6 +477,7 @@ def sparse_gqa_packed_decode_triton(q, k, v, indices, cu_q, cu_k, kv_lens, scale k, v, out, + out, indices, cu_q, cu_k, @@ -471,6 +496,8 @@ def sparse_gqa_packed_decode_triton(q, k, v, indices, cu_q, cu_k, kv_lens, scale out.stride(0), out.stride(1), out.stride(2), + 0, + 0, indices.stride(0), indices.stride(1) if indices.ndim == 3 else 0, indices.stride(2) if indices.ndim == 3 else indices.stride(1), @@ -479,6 +506,7 @@ def sparse_gqa_packed_decode_triton(q, k, v, indices, cu_q, cu_k, kv_lens, scale BLOCK_M=block_m, BLOCK_N=block_n, HEAD_DIM=head_dim, + RETURN_LSE=False, SLOT_INDICES=False, num_warps=warps, num_stages=stages, @@ -565,6 +593,7 @@ def sparse_gqa_paged_prefill_triton( k_cache, v_cache, out, + out, slots, cu_q, cu_q, @@ -583,6 +612,8 @@ def sparse_gqa_paged_prefill_triton( out.stride(0), out.stride(1), out.stride(2), + 0, + 0, slots.stride(0), 0, slots.stride(1), @@ -591,6 +622,7 @@ def sparse_gqa_paged_prefill_triton( BLOCK_M=block_m, BLOCK_N=block_n, HEAD_DIM=head_dim, + RETURN_LSE=False, SLOT_INDICES=True, num_warps=warps, num_stages=stages, @@ -598,6 +630,68 @@ def sparse_gqa_paged_prefill_triton( return out +def sparse_gqa_packed_decode_with_lse_triton( + q, k, v, indices, scale +) -> tuple[torch.Tensor, torch.Tensor]: + """Owner-local sparse decode returning FP32 normalized O and natural-log LSE.""" + + if not q.is_cuda: + raise ValueError("production owner sparse attention requires CUDA tensors") + k, v, indices = k.contiguous(), v.contiguous(), indices.contiguous() + total_q, num_q_heads, head_dim = q.shape + num_kv_heads = k.shape[1] + group_size = num_q_heads // num_kv_heads + block_m = max(16, triton.next_power_of_2(group_size)) + block_n, warps, stages = _get_best_config(total_q) + out = torch.empty(q.shape, dtype=torch.float32, device=q.device) + lse = torch.empty(q.shape[:2], dtype=torch.float32, device=q.device) + cu_q = torch.arange(total_q + 1, dtype=torch.int32, device=q.device) + cu_k = torch.zeros(total_q + 1, dtype=torch.int32, device=q.device) + kv_lens = torch.full( + (total_q,), indices.shape[-1], dtype=torch.int32, device=q.device + ) + _sparse_gqa_chunk_prefill[(1, total_q * num_kv_heads)]( + q, + k, + v, + out, + lse, + indices, + cu_q, + cu_k, + kv_lens, + scale, + indices.shape[-1], + q.stride(0), + q.stride(1), + q.stride(2), + k.stride(0), + k.stride(1), + k.stride(2), + v.stride(0), + v.stride(1), + v.stride(2), + out.stride(0), + out.stride(1), + out.stride(2), + lse.stride(0), + lse.stride(1), + indices.stride(0), + indices.stride(1) if indices.ndim == 3 else 0, + indices.stride(2) if indices.ndim == 3 else indices.stride(1), + NUM_KV_HEADS=num_kv_heads, + GROUP_SIZE=group_size, + BLOCK_M=block_m, + BLOCK_N=block_n, + HEAD_DIM=head_dim, + RETURN_LSE=True, + SLOT_INDICES=False, + num_warps=warps, + num_stages=stages, + ) + return out, lse + + @triton.jit def _fa2_valid_counts( seq_lens, @@ -802,5 +896,6 @@ def qwen_sparse_kv_extraction_compact_triton( "sparse_gqa_fwd_interface_triton", "sparse_gqa_fwd_interface_triton_ck", "sparse_gqa_packed_decode_triton", + "sparse_gqa_packed_decode_with_lse_triton", "sparse_gqa_paged_prefill_triton", ] diff --git a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py index b95ffe238dc8..d9175e007dc2 100644 --- a/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py +++ b/python/sglang/srt/layers/attention/qwen_sparse_attn_backend.py @@ -10,13 +10,17 @@ import math from copy import copy from functools import lru_cache -from typing import Dict, Optional, Tuple +from typing import Optional, Tuple import msgspec import torch import torch.nn.functional as F from sglang.srt.layers.attention.base_attn_backend import AttentionBackend +from sglang.srt.layers.attention.qsa.cache_sharding import ( + assert_qsa_cache_sharding_runtime_match, + get_qsa_cache_sharding_runtime, +) from sglang.srt.layers.attention.qsa.config import ( is_qwen_qsa, parse_qsa_profile, @@ -30,8 +34,14 @@ QSAIndexerMetadata, build_group_ring_slots, build_pending_ring_slots, + build_qsa_capture_token_slot_table, build_rope_position_matrix, compressed_decode_view, + localize_compressed_page_table, +) +from sglang.srt.layers.attention.qsa.owner_sparse_attn import ( + owner_sparse_attention, + select_kv_heads_for_query_shard, ) from sglang.srt.layers.attention.qsa.sparse_attn import ( pack_qsa_prefill_kv, @@ -53,7 +63,7 @@ from sglang.srt.mem_cache.memory_pool import HybridLinearKVPool, MHATokenToKVPool from sglang.srt.mem_cache.qsa_kv_pool import QSATokenToKVPool from sglang.srt.model_executor.forward_mode import ForwardMode -from sglang.srt.runtime_context import get_spec, get_stream +from sglang.srt.runtime_context import get_flags, get_parallel, get_spec, get_stream from sglang.srt.utils import is_hip, is_sm100_supported, is_sm120 logger = logging.getLogger(__name__) @@ -149,11 +159,11 @@ class QwenSparseAttnMetadata(msgspec.Struct, frozen=True): token_to_batch_idx: torch.Tensor token_slot_table: torch.Tensor indexer_metadata: QSAIndexerMetadata - row_req_pool_indices: Optional[torch.Tensor] = None + row_req_pool_indices: torch.Tensor | None = None is_cuda_graph: bool = False - fa2_valid_counts: Optional[torch.Tensor] = None - fa2_cu_seqlens_k: Optional[torch.Tensor] = None - fa2_cu_seqlens_q: Optional[torch.Tensor] = None + fa2_valid_counts: torch.Tensor | None = None + fa2_cu_seqlens_k: torch.Tensor | None = None + fa2_cu_seqlens_q: torch.Tensor | None = None class QSAMTPSharedSparseIndices: @@ -233,6 +243,24 @@ def __init__(self, runner=None) -> None: config = getattr(model_config, "hf_text_config", None) if config is None: config = getattr(model_config, "hf_config", None) + cache_sharding_runtime = get_qsa_cache_sharding_runtime() + if self.token_to_kv_pool is None: + pool_runtime = cache_sharding_runtime + else: + pool_runtime = getattr( + self.token_to_kv_pool, "cache_sharding_runtime", None + ) + if pool_runtime is None: + raise ValueError( + "backend QSA cache-sharding runtime mismatch: QSA pool did " + "not publish its construction contract" + ) + assert_qsa_cache_sharding_runtime_match( + cache_sharding_runtime, + pool_runtime, + component="backend", + ) + self.cache_sharding_runtime = cache_sharding_runtime self.qsa_profile = parse_qsa_profile(config) self.max_context_len = int(getattr(model_config, "context_len", 0)) self.compress_ratio = ( @@ -243,14 +271,14 @@ def __init__(self, runner=None) -> None: req_pool = getattr(runner, "req_to_token_pool", None) self.req_to_token = getattr(req_pool, "req_to_token", None) self.req_to_token_pool = req_pool - self.forward_metadata: Optional[QwenSparseAttnMetadata] = None - self._cuda_graph_metadata: Dict[ - Tuple[ForwardMode, int], QwenSparseAttnMetadata + self.forward_metadata: QwenSparseAttnMetadata | None = None + self._cuda_graph_metadata: dict[ + tuple[ForwardMode, int], QwenSparseAttnMetadata ] = {} self._cuda_graph_max_tokens = 0 - self._fa2_scratch: Dict[ - Tuple[int, int, torch.dtype, torch.device], - Tuple[torch.Tensor, torch.Tensor], + self._fa2_scratch: dict[ + tuple[int, int, torch.dtype, torch.device], + tuple[torch.Tensor, torch.Tensor], ] = {} self._graph_seq_lens = None self._graph_token_to_batch = None @@ -258,8 +286,10 @@ def __init__(self, runner=None) -> None: self._graph_fa2_valid_counts = None self._graph_fa2_cu_seqlens_k = None self._graph_write_locs = None + self._graph_write_owner_mask = None self._graph_compressed_page_table = None self._graph_compressed_lengths = None + self._graph_block_positions = None self._graph_prefix_lengths = None self._graph_dummy_token_slot_table = None self._graph_dummy_out_cache_loc = None @@ -271,6 +301,7 @@ def __init__(self, runner=None) -> None: self._graph_extend_lens = None self._graph_extend_lens_pin = None self._verify_mask = None + self._qsa_cache_sharding_buffers = None # (layer_id, k_full, v_full) of the prefill-CP K/V gather prefetch_cp_kv # started; that layer's _forward_extend_cp joins it. self._cp_kv_pending: Optional[Tuple[int, torch.Tensor, torch.Tensor]] = None @@ -281,6 +312,84 @@ def _is_speculative_paged_mode(forward_mode) -> bool: return False return forward_mode.is_target_verify() or forward_mode.is_draft_extend_v2() + @staticmethod + def _select_attention_kv_heads( + k: torch.Tensor, v: torch.Tensor, layer + ) -> tuple[torch.Tensor, torch.Tensor]: + parallel = get_parallel() + total_query_heads = layer.tp_q_head_num * parallel.attn_tp_size + total_kv_heads = getattr( + getattr(layer, "config", None), + "num_key_value_heads", + k.shape[1], + ) + return select_kv_heads_for_query_shard( + k, + v, + total_query_heads=total_query_heads, + total_kv_heads=total_kv_heads, + tp_size=parallel.attn_tp_size, + tp_rank=parallel.attn_tp_rank, + ) + + def _owner_sparse_attn_components(self): + """Resolve owner communication after the KV pool has finished setup.""" + + pool = self.token_to_kv_pool + group = getattr(pool, "qsa_owner_group", None) + runtime = getattr(pool, "qsa_raw_kv_sharding", None) + if (group is None) != (runtime is None): + raise ValueError( + "QSA owner group and raw-KV runtime must be initialized together" + ) + return group, runtime + + def _store_kv_cache(self, layer, loc, k, v) -> None: + pool = self.token_to_kv_pool + if getattr(pool, "qsa_raw_kv_sharding", None) is not None: + pool.set_kv_buffer(layer, loc, k, v) + return + self.kv_layers.layer(layer.layer_id).store(layer, loc, k, v) + + def _dispatch_owner_sparse_attention( + self, + q: torch.Tensor, + layer, + forward_batch, + topk_indices: torch.Tensor, + ) -> torch.Tensor | None: + group, runtime = self._owner_sparse_attn_components() + if group is None and runtime is None: + return None + if group is None or runtime is None: + raise ValueError( + "owner_sparse_attn_group and owner_sparse_attn_runtime " + "must be supplied together" + ) + metadata = self._resolve_metadata(forward_batch) + topk_indices = self._expand_block_indices(topk_indices, metadata) + global_token_slots = self._logical_to_physical(topk_indices, metadata) + pool = self.token_to_kv_pool + k_cache, v_cache = self._select_attention_kv_heads( + pool.get_key_buffer(layer.layer_id), + pool.get_value_buffer(layer.layer_id), + layer, + ) + static_buffers = self._qsa_cache_sharding_buffers + if static_buffers is not None and q.shape[0] > static_buffers.max_rows: + static_buffers = None + output, _ = owner_sparse_attention( + q, + k_cache, + v_cache, + global_token_slots, + runtime, + group=group, + softmax_scale=layer.scaling, + static_buffers=static_buffers, + ) + return output.to(q.dtype) + def _require_chain_speculation(self, forward_mode, spec_info) -> None: if forward_mode is None or not forward_mode.is_target_verify(): return @@ -391,7 +500,7 @@ def _graph_speculative_layout( forward_mode, spec_info, num_padding: int = 0, - ) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: base_lengths = cls._as_cpu_int_tensor(seq_lens_cpu, bs) if forward_mode.is_target_verify(): extend_len = int(spec_info.draft_token_num) @@ -515,7 +624,8 @@ def _qsa_write_plan( prefix_lens=None, ): """Compact per-row block ranges into ``capacity`` write entries, - a shape-derived bound (no sync); padding writes the inert reserved slot 0.""" + a shape-derived bound (no sync); ``valid`` distinguishes real groups + from padding instead of relying on the reserved slot value.""" device = token_slot_table.device # The table width is a host-side bound; # assert on device so a short table fails loudly without a sync. @@ -546,7 +656,7 @@ def _qsa_write_plan( row_token_starts[rows] + blocks * compress_ratio - prefix_lens[rows], torch.zeros_like(blocks), ) - return write_locs, group_end_positions, rows, member_rows + return write_locs, group_end_positions, rows, member_rows, valid def _qsa_build_write_plan( self, @@ -690,20 +800,40 @@ def _metadata_from_forward_batch(self, forward_batch) -> QwenSparseAttnMetadata: prefill_sequence_lengths_cpu = tuple( forward_batch.seq_lens_cpu[:batch_size].tolist() ) + if get_flags().dp.capturing_prefill_graph: + token_slot_table = build_qsa_capture_token_slot_table( + batch_size=sequence_lengths.numel(), + max_length=max_length, + device=sequence_lengths.device, + ) decode_page_table = None decode_lengths = None + decode_block_positions = None decode_logical_positions = None pending_ring_slots = None compress_group_ring_locs = None extend_rope_matrix = None - write_locs, group_positions, group_sequence_ids, group_member_rows = ( - self._qsa_build_write_plan( - forward_batch=forward_batch, - speculative_paged=speculative_paged, - token_slot_table=token_slot_table, - sequence_lengths=sequence_lengths, - ) + ( + write_locs, + group_positions, + group_sequence_ids, + group_member_rows, + write_valid_mask, + ) = self._qsa_build_write_plan( + forward_batch=forward_batch, + speculative_paged=speculative_paged, + token_slot_table=token_slot_table, + sequence_lengths=sequence_lengths, ) + pool = self.token_to_kv_pool + sharding = getattr(pool, "qsa_compressed_sharding", None) + if sharding is None: + write_owner_mask = write_valid_mask + local_write_locs = torch.where(write_valid_mask, write_locs.long(), 0) + else: + write_owner_mask, local_write_locs = sharding.local_write_targets( + write_locs.long(), valid_mask=write_valid_mask + ) decode_like = speculative_paged or forward_batch.forward_mode.is_decode() if decode_like: decode_logical_positions = ( @@ -719,15 +849,70 @@ def _metadata_from_forward_batch(self, forward_batch) -> QwenSparseAttnMetadata: ring_logical_positions = extend_positions.flatten()[ : token_to_batch_idx.numel() ] + prefill_page_table = None + prefill_lengths = None + prefill_block_positions = None + prefill_cache_indices = None + prefill_sequence_lengths_cpu = None + if not decode_like and forward_batch.seq_lens_cpu is not None: + host_sequence_lengths = forward_batch.seq_lens_cpu[ + : sequence_lengths.numel() + ] + prefill_sequence_lengths_cpu = tuple( + int(length) for length in host_sequence_lengths.tolist() + ) if not self.should_reuse_mtp_sparse_indices(forward_batch): - if decode_like: - pool = self.token_to_kv_pool - decode_page_table, decode_lengths = compressed_decode_view( - compressed_page_size=pool.qsa_compressed_page_size, - compress_ratio=pool.qsa_compress_ratio, - sequence_lengths=sequence_lengths, - token_slot_table=token_slot_table, + global_page_table, global_lengths = compressed_decode_view( + compressed_page_size=pool.qsa_compressed_page_size, + compress_ratio=pool.qsa_compress_ratio, + sequence_lengths=sequence_lengths, + token_slot_table=token_slot_table, + ) + if sharding is not None: + ( + local_page_table, + local_lengths, + local_block_positions, + ) = localize_compressed_page_table( + global_page_table=global_page_table, + compressed_lengths=global_lengths, + sharding=sharding, ) + if decode_like: + ( + decode_page_table, + decode_lengths, + decode_block_positions, + ) = (local_page_table, local_lengths, local_block_positions) + else: + ( + prefill_page_table, + prefill_lengths, + prefill_block_positions, + ) = ( + local_page_table, + local_lengths, + local_block_positions, + ) + page_size = pool.qsa_compressed_page_size + local_block_table = ( + local_page_table.long().clamp_min(0).unsqueeze(-1) * page_size + + torch.arange( + page_size, + device=local_page_table.device, + dtype=torch.long, + ).view(1, 1, -1) + ).flatten(1) + valid_local_blocks = torch.arange( + local_block_table.shape[1], + device=local_block_table.device, + ).unsqueeze(0) < local_lengths.unsqueeze(1) + prefill_cache_indices = local_block_table[ + valid_local_blocks + ].contiguous() + elif decode_like: + decode_page_table = global_page_table + decode_lengths = global_lengths pending_ring_slots = build_pending_ring_slots( token_to_batch_idx=token_to_batch_idx, req_pool_indices=row_req_pool_indices, @@ -763,11 +948,18 @@ def _metadata_from_forward_batch(self, forward_batch) -> QwenSparseAttnMetadata: block_topk=self.token_to_kv_pool.qsa_block_topk, req_pool_indices=row_req_pool_indices, write_locs=write_locs, + write_owner_mask=write_owner_mask, + local_write_locs=local_write_locs, compress_group_positions=group_positions, compress_sequence_ids=group_sequence_ids, compress_member_rows=group_member_rows, + prefill_page_table=prefill_page_table, + prefill_lengths=prefill_lengths, + prefill_block_positions=prefill_block_positions, + prefill_cache_indices=prefill_cache_indices, decode_page_table=decode_page_table, decode_lengths=decode_lengths, + decode_block_positions=decode_block_positions, decode_logical_positions=decode_logical_positions, pending_ring_slots=pending_ring_slots, compress_group_ring_locs=compress_group_ring_locs, @@ -829,7 +1021,7 @@ def init_forward_metadata_out_graph( ) @property - def verify_mask(self) -> Optional[VerifyMask]: + def verify_mask(self) -> VerifyMask | None: return self._verify_mask def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: @@ -847,6 +1039,35 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: is_read=False, ) self._cuda_graph_max_tokens = max_num_tokens + owner_group, _ = self._owner_sparse_attn_components() + if owner_group is not None: + from sglang.srt.layers.attention.qsa.cache_sharding import ( + QSACacheShardingStaticBuffers, + ) + + config = self.runner.model_config.hf_text_config + parallel = get_parallel() + gathered_query_heads = ( + int(config.num_attention_heads) + // int(parallel.attn_tp_size) + * int(owner_group.world_size) + ) + self._qsa_cache_sharding_buffers = QSACacheShardingStaticBuffers( + max_rows=max_num_tokens, + topk=max(1, self.qsa_profile.block_topk), + owner_topk=( + self.qsa_profile.budget + self.qsa_profile.compress_ratio - 1 + ), + num_heads=gathered_query_heads, + head_dim=int(config.head_dim), + world_size=int(owner_group.world_size), + device=self.device, + dtype=self.runner.dtype, + candidate_group=owner_group, + ) + self.token_to_kv_pool.qsa_cache_sharding_buffers = ( + self._qsa_cache_sharding_buffers + ) max_blocks = math.ceil(self.max_context_len / self.compress_ratio) max_pages = max( 1, @@ -870,12 +1091,24 @@ def init_cuda_graph_state(self, max_bs: int, max_num_tokens: int) -> None: self._graph_write_locs = torch.zeros( max_num_tokens, dtype=torch.int32, device=self.device ) + self._graph_write_owner_mask = torch.zeros( + max_num_tokens, dtype=torch.bool, device=self.device + ) self._graph_compressed_page_table = torch.zeros( (max_num_tokens, max_pages), dtype=torch.int32, device=self.device ) self._graph_compressed_lengths = torch.zeros( max_num_tokens, dtype=torch.int32, device=self.device ) + self._graph_block_positions = torch.full( + ( + max_num_tokens, + max_pages * self.token_to_kv_pool.qsa_compressed_page_size, + ), + -1, + dtype=torch.int32, + device=self.device, + ) self._graph_prefix_lengths = torch.zeros( max_num_tokens, dtype=torch.int32, device=self.device ) @@ -974,10 +1207,12 @@ def _capture_cuda_graph_metadata( req_pool_indices=self._graph_row_req_pool_indices[:metadata_rows], is_cuda_graph=True, graph_write_locs=self._graph_write_locs[:metadata_rows], + write_owner_mask=self._graph_write_owner_mask[:metadata_rows], graph_compressed_page_table=self._graph_compressed_page_table[ :metadata_rows ], graph_compressed_lengths=self._graph_compressed_lengths[:metadata_rows], + decode_block_positions=self._graph_block_positions[:metadata_rows], graph_prefix_lengths=self._graph_prefix_lengths[:metadata_rows], decode_logical_positions=self._graph_logical_positions[:metadata_rows], pending_ring_slots=self._graph_state_slots[:metadata_rows], @@ -1331,9 +1566,8 @@ def _resolve_metadata(self, forward_batch) -> QwenSparseAttnMetadata: assert self.forward_metadata is not None return self.forward_metadata - @staticmethod def _logical_to_physical( - logical_indices: torch.Tensor, metadata: QwenSparseAttnMetadata + self, logical_indices: torch.Tensor, metadata: QwenSparseAttnMetadata ) -> torch.Tensor: sequence_ids = metadata.token_to_batch_idx.long() if sequence_ids.numel() != logical_indices.shape[0]: @@ -1342,14 +1576,30 @@ def _logical_to_physical( 0, sequence_ids ) valid = (logical_indices >= 0) & (logical_indices < row_lengths.unsqueeze(1)) - safe = logical_indices.clamp( - min=0, max=metadata.token_slot_table.shape[1] - 1 - ).long() - slots = metadata.token_slot_table[sequence_ids[:, None], safe] - return torch.where(valid, slots, torch.full_like(slots, -1)).to(torch.int32) + if metadata.is_cuda_graph: + if self.req_to_token is None or metadata.row_req_pool_indices is None: + raise RuntimeError( + "QSA graph replay requires live req_to_token row metadata" + ) + safe = logical_indices.clamp( + min=0, max=self.req_to_token.shape[1] - 1 + ).long() + reqs = metadata.row_req_pool_indices.long().index_select(0, sequence_ids) + slots = self.req_to_token[reqs[:, None], safe] + else: + safe = logical_indices.clamp( + min=0, max=metadata.token_slot_table.shape[1] - 1 + ).long() + slots = metadata.token_slot_table[sequence_ids[:, None], safe] + physical_slots = torch.where(valid, slots, torch.full_like(slots, -1)).to( + torch.int32 + ) + return physical_slots @staticmethod def _supports_fused_kv_pool(pool): + if getattr(getattr(pool, "cache_sharding_runtime", None), "enabled", False): + return False if type(pool) in (HybridLinearKVPool, QSATokenToKVPool): storage = pool.full_kv_pool elif type(pool) is MHATokenToKVPool: @@ -1464,7 +1714,7 @@ def forward_extend( layer, forward_batch, save_kv_cache: bool = True, - topk_indices: Optional[torch.Tensor] = None, + topk_indices: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: if topk_indices is None: @@ -1481,9 +1731,7 @@ def forward_extend( ) if fused_output is not None: return fused_output - self.kv_layers.layer(layer.layer_id).store( - layer, forward_batch.out_cache_loc, k, v - ) + self._store_kv_cache(layer, forward_batch.out_cache_loc, k, v) q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) num_output_rows = q.shape[0] num_valid_rows = topk_indices.shape[0] @@ -1500,6 +1748,11 @@ def forward_extend( q, layer, forward_batch, topk_indices ) return self._pad_extend_output(output, num_output_rows) + output = self._dispatch_owner_sparse_attention( + q, layer, forward_batch, topk_indices + ) + if output is not None: + return self._pad_extend_output(output, num_output_rows) if not q.is_cuda: metadata = self._resolve_metadata(forward_batch) slots = self._logical_to_physical(topk_indices, metadata) @@ -1619,9 +1872,7 @@ def _forward_extend_cp( if save_kv_cache: # out_cache_loc was trimmed to the full extend token count by # prepare_cp_forward; write all new tokens on every rank. - self.kv_layers.layer(layer.layer_id).store( - layer, forward_batch.out_cache_loc, k_full, v_full - ) + self._store_kv_cache(layer, forward_batch.out_cache_loc, k_full, v_full) num_local_rows = int(meta.total_q_prev_tokens + meta.total_q_next_tokens) q = q[:num_local_rows].contiguous() indices = topk_indices[:num_local_rows].to(torch.int32).contiguous() @@ -1720,7 +1971,7 @@ def _get_fa2_scratch( head_dim: int, dtype: torch.dtype, device: torch.device, - ) -> Tuple[torch.Tensor, torch.Tensor]: + ) -> tuple[torch.Tensor, torch.Tensor]: key = (num_kv_heads, head_dim, dtype, device) buffers = self._fa2_scratch.get(key) if buffers is None or buffers[0].shape[0] < capacity: @@ -1897,7 +2148,7 @@ def forward_decode( layer, forward_batch, save_kv_cache: bool = True, - topk_indices: Optional[torch.Tensor] = None, + topk_indices: torch.Tensor | None = None, **kwargs, ) -> torch.Tensor: if topk_indices is None: @@ -1908,11 +2159,12 @@ def forward_decode( ) if fused_output is not None: return fused_output - self.kv_layers.layer(layer.layer_id).store( - layer, forward_batch.out_cache_loc, k, v - ) + self._store_kv_cache(layer, forward_batch.out_cache_loc, k, v) q = q.reshape(-1, layer.tp_q_head_num, layer.head_dim) - return self._forward_paged_attention(q, layer, forward_batch, topk_indices) + output = self._forward_paged_attention( + q, layer, forward_batch, topk_indices + ).reshape(q.shape[0], -1) + return output def _forward_paged_attention( self, @@ -1921,6 +2173,11 @@ def _forward_paged_attention( forward_batch, topk_indices: torch.Tensor, ) -> torch.Tensor: + output = self._dispatch_owner_sparse_attention( + q, layer, forward_batch, topk_indices + ) + if output is not None: + return output.reshape(q.shape[0], -1) pool = self.token_to_kv_pool k_buffer = pool.layer_router.read(layer.layer_id).key() v_buffer = pool.layer_router.read(layer.layer_id).value() @@ -2233,8 +2490,8 @@ def init_forward_metadata_in_graph(self, forward_batch) -> None: __all__ = [ - "is_qwen_qsa", - "QwenSparseAttnMetadata", "QwenSparseAttnBackend", + "QwenSparseAttnMetadata", "QwenSparseMultiStepDraftBackend", + "is_qwen_qsa", ] diff --git a/python/sglang/srt/managers/scheduler.py b/python/sglang/srt/managers/scheduler.py index e9c57479c2fd..999471080750 100644 --- a/python/sglang/srt/managers/scheduler.py +++ b/python/sglang/srt/managers/scheduler.py @@ -2635,7 +2635,7 @@ def init_dp_attn_adapter(self) -> None: ) def init_pool_stats_observer(self) -> None: - global_tokens = self.max_total_num_tokens * self.kv_shard_widening + global_tokens = self.max_total_num_tokens if is_cp_sharded_padded(): global_tokens = from_cp_sharded(global_tokens) storage = kv_storage_of( @@ -2727,13 +2727,13 @@ def init_load_inquirer(self) -> None: maxlen=PAUSE_WINDOWS_MAX ) self._pause_started_at: Optional[float] = None - global_tokens = self.max_total_num_tokens * self.kv_shard_widening + global_tokens = self.max_total_num_tokens if is_cp_sharded_padded(): global_tokens = from_cp_sharded(global_tokens) self.load_inquirer = SchedulerLoadInquirer( disaggregation_mode=self.disaggregation_mode, server_args=self.server_args, - # The worker reports logical DCP capacity; add KV shard widening. + # The worker already reports logical DCP/page-interleave capacity. max_total_num_tokens=global_tokens, max_running_requests=self.max_running_requests, pool_stats_observer=self.pool_stats_observer, @@ -2887,9 +2887,8 @@ def init_req_max_new_tokens(self, req): ) # PrefillAdder reserves one page per shard; the allocator reserves one. # Subtract the other N - 1 pages to keep queue admission schedulable. - token_capacity = ( - self.max_total_num_tokens * self.kv_shard_widening - - self.page_size * (self.kv_shard_widening - 1) + token_capacity = self.max_total_num_tokens - self.page_size * ( + self.kv_shard_widening - 1 ) if is_cp_sharded_padded(): token_capacity = from_cp_sharded(token_capacity) diff --git a/python/sglang/srt/managers/tp_worker.py b/python/sglang/srt/managers/tp_worker.py index 0d9d925de9b8..708f8e3280d5 100644 --- a/python/sglang/srt/managers/tp_worker.py +++ b/python/sglang/srt/managers/tp_worker.py @@ -45,7 +45,6 @@ from sglang.srt.managers.schedule_batch import ScheduleBatch from sglang.srt.managers.utils import GenerationBatchResult from sglang.srt.mem_cache.allocator import BaseTokenToKVPoolAllocator -from sglang.srt.mem_cache.allocator.page_interleave import page_interleave_shard_size from sglang.srt.mem_cache.memory_pool import ReqToTokenPool from sglang.srt.model_executor.forward_batch_info import ForwardBatch, PPProxyTensors from sglang.srt.model_executor.forward_mode import CaptureHiddenMode @@ -548,10 +547,7 @@ def register_hisparse_coordinator(self, coordinator): def get_worker_info(self): # The runner already reports logical DCP capacity. - kv_capacity = ( - self.model_runner.effective_logical_max_total_num_tokens - * page_interleave_shard_size(self.model_runner.token_to_kv_pool_allocator) - ) + kv_capacity = self.model_runner.effective_logical_max_total_num_tokens if is_cp_sharded_padded(): kv_capacity = from_cp_sharded(kv_capacity) # Not an error: a deployment that never sends max-length requests may diff --git a/python/sglang/srt/mem_cache/allocator/page_interleave.py b/python/sglang/srt/mem_cache/allocator/page_interleave.py index 3bcab33f5014..73a6e6ed9b53 100644 --- a/python/sglang/srt/mem_cache/allocator/page_interleave.py +++ b/python/sglang/srt/mem_cache/allocator/page_interleave.py @@ -328,11 +328,51 @@ def alloc_decode( seq_lens: torch.Tensor, seq_lens_cpu: torch.Tensor, last_loc: torch.Tensor, + rotation_bases: Optional[List[int]] = None, ): - raise NotImplementedError( - "PageInterleavePoolAllocator does not support decode allocation " - "(logical-page KV sharding runs on prefill nodes only)" + ps, shard_size = self.page_size, self.shard_size + bs = len(seq_lens_cpu) + assert rotation_bases is not None and len(rotation_bases) == bs, ( + "sharded alloc_decode needs the rotation base of every request" ) + if self.debug_mode: + assert torch.all((last_loc + 2) % ps == seq_lens % ps), ( + "decode length and last location disagree" + ) + + # A decode token needs a fresh page exactly at the first position of a + # position-page. Its owner is host-computable from the chain rotation, + # so the hot path does not synchronize on last_loc. + new_page_reqs: List[List[int]] = [[] for _ in range(shard_size)] + for i in range(bs): + seq_len = int(seq_lens_cpu[i]) + if seq_len % ps != 1: + continue + owner = (int(rotation_bases[i]) + (seq_len - 1) // ps) % shard_size + new_page_reqs[owner].append(i) + + needs = [len(reqs) for reqs in new_page_reqs] + available = self.class_free_page_counts() + if any(need > available[c] for c, need in enumerate(needs)): + return None + if self.need_sort and any( + needs[c] > len(self.class_free_pages[c]) for c in range(shard_size) + ): + self.merge_and_sort_free() + + out_indices = last_loc + 1 + for owner, reqs in enumerate(new_page_reqs): + count = len(reqs) + if count == 0: + continue + pages = self.class_free_pages[owner][:count] + self.class_free_pages[owner] = self.class_free_pages[owner][count:] + request_indices = torch.tensor(reqs, dtype=torch.int64, device=self.device) + out_indices[request_indices] = pages.to(dtype=out_indices.dtype) * ps + + if self.debug_mode: + assert len(torch.unique(out_indices)) == len(out_indices) + return out_indices def free(self, free_index: torch.Tensor): if free_index.numel() == 0: diff --git a/python/sglang/srt/mem_cache/kv_cache_configurator.py b/python/sglang/srt/mem_cache/kv_cache_configurator.py index 1e32cbd6c976..b7d353b7ac2f 100644 --- a/python/sglang/srt/mem_cache/kv_cache_configurator.py +++ b/python/sglang/srt/mem_cache/kv_cache_configurator.py @@ -112,6 +112,7 @@ get_schedule, get_spec, mamba_cache_chunk_size, + max_prefill_buffer_tokens, max_speculative_num_draft_tokens, pre_capture_activation_reserve_mb, ) @@ -130,6 +131,35 @@ logger = logging.getLogger(__name__) +def _qsa_cache_sharding_graph_reservation_gb( + *, + sharding_size: int, + is_qsa: bool, + cuda_graph_enabled: bool, + prefill_tokens: int, + hidden_size: int, + hyper_connection_count: int, + activation_element_size: int, +) -> float: + if sharding_size <= 1 or not is_qsa or not cuda_graph_enabled: + return 0.0 + # NCCL graph-private pools are created per QSA cache shard during decode + # capture and are not visible when the KV pool is initially sized. Qwen4 + # PLE holds the indexed short-conv result while SiLU materializes its output, + # then the enclosing residual add materializes another result. Thus three + # [prefill_tokens, hc_count * hidden_size] tensors can overlap after those + # pools exist; keep them out of the KV budget too. + graph_private_gb = sharding_size * 0.5 + ple_selection_gb = ( + prefill_tokens + * hidden_size + * hyper_connection_count + * activation_element_size + / (1 << 30) + ) + return graph_private_gb + 3 * ple_selection_gb + + def _should_elide_dsa_index_k(*, is_draft_worker: bool) -> bool: """Shared-topk layers get a 0-row index-K buffer; HiCache and PD must skip those entries.""" memory_config = get_memory() @@ -473,8 +503,15 @@ def hybrid_swa_token_capacity( ) return full_capacity or swa_capacity - def logical_token_capacity(self, *, max_total_num_tokens: int) -> int: + def logical_token_capacity( + self, + *, + max_total_num_tokens: int, + allocator: Optional[BaseTokenToKVPoolAllocator] = None, + ) -> int: """Request tokens a pool of `max_total_num_tokens` per-rank rows holds.""" + if isinstance(allocator, PageInterleavePoolAllocator): + return allocator.size # SWA allocators never widen under DCP. if self.is_hybrid_swa: return max_total_num_tokens @@ -2550,6 +2587,9 @@ def _build_hybrid_linear_kv_pool( mla_pool_class=mla_pool_class, dsa_pool_class=dsa_pool_class, ) + from sglang.srt.layers.attention.qsa.cache_sharding import ( + get_qsa_cache_sharding_runtime, + ) from sglang.srt.layers.attention.qsa.config import ( parse_qsa_profile, ) @@ -2577,6 +2617,7 @@ def _build_hybrid_linear_kv_pool( qsa_token_topk=qsa_profile.budget, num_request_slots=req_to_token_pool.req_to_token.shape[0], qsa_indexer_dtype=resolve_qsa_indexer_dtype(qsa_indexer_dtype), + cache_sharding_runtime=get_qsa_cache_sharding_runtime(), ) if self.modelslim_cache_layout is not None: full_pool_class = partial( @@ -2698,7 +2739,26 @@ def _build_token_to_kv_pool_allocator( # Initialize token_to_kv_pool_allocator need_sort = get_disagg().disaggregation_mode in ("decode", "prefill") if token_to_kv_pool_allocator is None: - if _is_npu and ( + from sglang.srt.mem_cache.qsa_kv_pool import QSATokenToKVPool + + qsa_sharding_runtime = getattr( + token_to_kv_pool, "cache_sharding_runtime", None + ) + if ( + isinstance(token_to_kv_pool, QSATokenToKVPool) + and qsa_sharding_runtime is not None + and qsa_sharding_runtime.enabled + ): + token_to_kv_pool_allocator = PageInterleavePoolAllocator( + size=sizes.max_total_num_tokens, + physical_page_size=get_schedule().page_size, + shard_size=qsa_sharding_runtime.size, + dtype=self.kv_cache_dtype, + device=self.device, + kvcache=token_to_kv_pool, + need_sort=need_sort, + ) + elif _is_npu and ( get_exec().kernel.attention_backend == "ascend" or is_dsv4_model or self.hybrid_gdn_config is not None @@ -2897,6 +2957,29 @@ def _profile_available_bytes(self, pre_model_load_memory: int) -> int: ) slack_gb = pre_model_load_memory * (1 - get_schedule().mem_fraction_static) + from sglang.srt.layers.attention.qsa.config import parse_qsa_profile + + qsa_profile = parse_qsa_profile(self.model_config.hf_config) + text_config = self.model_config.hf_text_config + qsa_sharding_graph_reservation_gb = _qsa_cache_sharding_graph_reservation_gb( + sharding_size=int(get_parallel().qsa_cache_sharding_size), + is_qsa=qsa_profile is not None, + cuda_graph_enabled=( + get_exec().graph.cuda_graph_config.decode.backend != Backend.DISABLED + ), + prefill_tokens=max_prefill_buffer_tokens(), + hidden_size=int(getattr(text_config, "hidden_size", 0)), + hyper_connection_count=int(getattr(text_config, "hc_count", 1)), + activation_element_size=self.model_dtype.itemsize, + ) + if qsa_sharding_graph_reservation_gb: + slack_gb += qsa_sharding_graph_reservation_gb + logger.info( + "Reserving %.2f GB for QSA cache-sharding CUDA-graph " + "collective private pools.", + qsa_sharding_graph_reservation_gb, + ) + # Every prefill batch fits a captured bucket when batches are capped at the # request slots; the exemption also needs an explicit Mamba size: a larger # rest_memory would otherwise auto-size a bigger physical Mamba pool below. diff --git a/python/sglang/srt/mem_cache/qsa_kv_pool.py b/python/sglang/srt/mem_cache/qsa_kv_pool.py index 77eee28ba844..c58096abe7cb 100644 --- a/python/sglang/srt/mem_cache/qsa_kv_pool.py +++ b/python/sglang/srt/mem_cache/qsa_kv_pool.py @@ -10,11 +10,19 @@ import logging from contextlib import nullcontext -from typing import List, Optional +from dataclasses import dataclass import torch from sglang.srt.constants import GPU_MEMORY_TYPE_KV_CACHE +from sglang.srt.layers.attention.qsa.cache_sharding import ( + QSACacheShardingRuntime, + assert_qsa_cache_sharding_runtime_match, + get_qsa_cache_sharding_runtime, +) +from sglang.srt.mem_cache.allocator.page_interleave import ( + page_interleave_shard_size, +) from sglang.srt.mem_cache.device_pool_info import ( DevicePoolInfo, IndexKeyBufferInfo, @@ -22,6 +30,10 @@ ) from sglang.srt.mem_cache.hicache_storage import PoolName from sglang.srt.mem_cache.memory_pool import GB, HybridLinearKVPool, MambaPool +from sglang.srt.mem_cache.page_interleave import ( + PageInterleavePlacement, + PageShardSpec, +) logger = logging.getLogger(__name__) @@ -51,6 +63,228 @@ def resolve_qsa_indexer_dtype(name: str) -> torch.dtype: ) +_INTEGER_DTYPES = ( + torch.uint8, + torch.int8, + torch.int16, + torch.int32, + torch.int64, +) + + +def assert_qsa_indices_in_bounds( + indices: torch.Tensor, + upper_bound: int, + *, + valid_mask: torch.Tensor | None = None, + label: str, +) -> None: + if not isinstance(indices, torch.Tensor): + raise TypeError(f"{label} must be a tensor") + if indices.dtype not in _INTEGER_DTYPES: + raise TypeError(f"{label} must have integer dtype") + if upper_bound < 0: + raise ValueError(f"{label} upper bound must be non-negative") + if valid_mask is None: + valid_mask = indices >= 0 + elif valid_mask.shape != indices.shape: + raise ValueError(f"{label} valid mask must match index shape") + else: + valid_mask = valid_mask.to(device=indices.device, dtype=torch.bool) + condition = torch.all((~valid_mask) | ((indices >= 0) & (indices < upper_bound))) + message = f"{label} must be in [0, {upper_bound})" + if indices.is_cuda: + torch._assert_async(condition, message) + elif not bool(condition.item()): + raise IndexError(message) + + +def _qsa_placement(*, rank: int, world_size: int, page_size: int): + if world_size <= 0 or not 0 <= rank < world_size: + raise ValueError(f"rank must be in [0, {world_size}), got {rank}") + return PageInterleavePlacement( + PageShardSpec( + shard_rank=rank, + shard_size=world_size, + page_size=page_size, + max_prefix_tokens=0, + chunk_tokens=0, + ) + ) + + +def qsa_global_compressed_capacity( + local_token_capacity: int, + *, + compress_ratio: int, + world_size: int | None = None, + allocator=None, +) -> int: + if allocator is not None: + allocator_world_size = page_interleave_shard_size(allocator) + if world_size is not None and world_size != allocator_world_size: + raise ValueError("allocator and explicit QSA shard widths disagree") + world_size = allocator_world_size + world_size = 1 if world_size is None else world_size + if min(local_token_capacity, compress_ratio, world_size) <= 0: + raise ValueError("QSA compressed-capacity inputs must be positive") + return -((local_token_capacity * world_size) // -compress_ratio) + + +@dataclass(frozen=True) +class QSARawKVSharding: + local_capacity: int + world_size: int + rank: int + + def __post_init__(self) -> None: + if self.local_capacity < 0: + raise ValueError("local_capacity must be non-negative") + object.__setattr__( + self, + "placement", + _qsa_placement(rank=self.rank, world_size=self.world_size, page_size=1), + ) + + @property + def global_capacity(self) -> int: + return self.local_capacity * self.world_size + + def global_to_local_tokens(self, global_slots: torch.Tensor) -> torch.Tensor: + global_slots = global_slots.long() + valid = (global_slots >= 0) & (global_slots < self.global_capacity) + owned = self.placement.local_mask(global_slots, self.rank) + local = self.placement.local_index(global_slots.clamp_min(0)) + return torch.where(valid & owned, local, -1) + + def local_write_targets( + self, global_slots: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + assert_qsa_indices_in_bounds( + global_slots, + self.global_capacity, + valid_mask=global_slots >= 0, + label="raw KV write locations", + ) + local = self.global_to_local_tokens(global_slots) + owner = local >= 0 + safe_local = torch.where(owner, local, 0) + assert_qsa_indices_in_bounds( + safe_local, + self.local_capacity, + valid_mask=owner, + label="raw KV local write locations", + ) + return owner, safe_local + + def local_copy_targets( + self, global_src: torch.Tensor, global_dst: torch.Tensor + ) -> tuple[torch.Tensor, torch.Tensor]: + if global_src.shape != global_dst.shape: + raise ValueError("source and destination slots must have matching shapes") + src = self.global_to_local_tokens(global_src) + dst = self.global_to_local_tokens(global_dst) + src_owned = src >= 0 + dst_owned = dst >= 0 + if bool(torch.any(src_owned != dst_owned).item()): + raise ValueError("cross-owner raw K/V copy requires communication") + return src[src_owned], dst[dst_owned] + + +@dataclass(frozen=True) +class QSACompressedBlockSharding: + """Interleaved page ownership; incomplete pending groups stay replicated.""" + + global_blocks: int + compressed_page_size: int + world_size: int + rank: int + + def __post_init__(self) -> None: + if self.global_blocks < 0: + raise ValueError("global_blocks must be non-negative") + if self.compressed_page_size <= 0: + raise ValueError( + "compressed_page_size must be positive, got " + f"{self.compressed_page_size}" + ) + object.__setattr__( + self, + "placement", + _qsa_placement( + rank=self.rank, + world_size=self.world_size, + page_size=self.compressed_page_size, + ), + ) + + @property + def global_pages(self) -> int: + return -(self.global_blocks // -self.compressed_page_size) + + @property + def local_pages(self) -> int: + if self.rank >= self.global_pages: + return 0 + return (self.global_pages - 1 - self.rank) // self.world_size + 1 + + @property + def local_blocks(self) -> int: + return self.local_pages * self.compressed_page_size + + def global_to_local_pages(self, global_pages: torch.Tensor) -> torch.Tensor: + global_pages = global_pages.long() + valid = (global_pages >= 0) & (global_pages < self.global_pages) + owned = global_pages.remainder(self.world_size) == self.rank + return torch.where( + valid & owned, + torch.div(global_pages, self.world_size, rounding_mode="floor"), + -1, + ) + + def global_to_local_blocks(self, global_blocks: torch.Tensor) -> torch.Tensor: + global_blocks = global_blocks.long() + valid = (global_blocks >= 0) & (global_blocks < self.global_blocks) + safe_blocks = global_blocks.clamp_min(0) + owned = self.placement.local_mask(safe_blocks, self.rank) + local_blocks = self.placement.local_index(safe_blocks) + return torch.where(valid & owned, local_blocks, -1) + + def local_write_targets( + self, + global_blocks: torch.Tensor, + *, + valid_mask: torch.Tensor | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if global_blocks.ndim != 1: + raise ValueError("global_blocks must be one-dimensional") + checked_mask = ( + global_blocks >= 0 + if valid_mask is None + else valid_mask.to(device=global_blocks.device, dtype=torch.bool) + ) + assert_qsa_indices_in_bounds( + global_blocks, + self.global_blocks, + valid_mask=checked_mask, + label="compressed KV write locations", + ) + local_blocks = self.global_to_local_blocks(global_blocks) + owner_mask = local_blocks >= 0 + if valid_mask is not None: + if valid_mask.shape != global_blocks.shape: + raise ValueError("valid_mask must match global_blocks") + owner_mask &= valid_mask.to(device=global_blocks.device, dtype=torch.bool) + safe_local_blocks = torch.where(owner_mask, local_blocks, 0) + assert_qsa_indices_in_bounds( + safe_local_blocks, + self.local_blocks, + valid_mask=owner_mask, + label="compressed KV local write locations", + ) + return owner_mask, safe_local_blocks + + class QSATokenToKVPool(HybridLinearKVPool): """Hybrid KV pool with the minimal BF16 state required by simple QSA.""" @@ -102,7 +336,7 @@ def __init__( page_size: int, head_num: int, head_dim: int, - full_attention_layer_ids: List[int], + full_attention_layer_ids: list[int], device: str, mamba_pool: MambaPool, qsa_index_kv_heads: int, @@ -110,10 +344,11 @@ def __init__( qsa_compress_ratio: int, qsa_token_topk: int, num_request_slots: int, + cache_sharding_runtime: QSACacheShardingRuntime | None = None, enable_memory_saver: bool = False, enable_kv_cache_copy: bool = False, - start_layer: Optional[int] = None, - full_kv_pool_class: Optional[type] = None, + start_layer: int | None = None, + full_kv_pool_class: type | None = None, quant_method=None, post_capture_active: bool = False, qsa_indexer_dtype: torch.dtype = torch.bfloat16, @@ -132,6 +367,15 @@ def __init__( self.qsa_key_state_buffer_pool = [] self.qsa_compressed_k_buffer_pool = [] self.qsa_rope_position_buffer = torch.empty(0) + topology_runtime = get_qsa_cache_sharding_runtime() + if cache_sharding_runtime is None: + cache_sharding_runtime = topology_runtime + assert_qsa_cache_sharding_runtime_match( + cache_sharding_runtime, + topology_runtime, + component="pool", + ) + self.cache_sharding_runtime = cache_sharding_runtime super().__init__( size=size, dtype=dtype, @@ -182,7 +426,14 @@ def __init__( # Compressed slots mirror the full-KV slot space 1:ratio; the "page" # seen by the scoring kernels is one full-KV page's worth of groups. self.qsa_compressed_page_size = page_size // self.qsa_compress_ratio - self.qsa_compressed_capacity = -(state_size // -self.qsa_compress_ratio) + self.qsa_global_compressed_capacity = qsa_global_compressed_capacity( + state_size, + compress_ratio=self.qsa_compress_ratio, + world_size=1, + ) + self.qsa_compressed_capacity = self.qsa_global_compressed_capacity + self.qsa_raw_kv_sharding = None + self.qsa_owner_group = None # Pre-compression index-K state is a per-request ring, not a per-token cache: # only the pending group's ``ratio`` members must survive a forward, # addressed as ``req_pool_idx * ratio + position % ratio``. @@ -193,6 +444,26 @@ def __init__( ) self.qsa_num_request_slots = int(num_request_slots) ring_slots = self.qsa_num_request_slots * self.qsa_compress_ratio + self.qsa_compressed_sharding = None + if cache_sharding_runtime.enabled: + self.qsa_global_compressed_capacity = qsa_global_compressed_capacity( + state_size, + compress_ratio=self.qsa_compress_ratio, + world_size=cache_sharding_runtime.size, + ) + self.qsa_owner_group = cache_sharding_runtime.group + self.qsa_raw_kv_sharding = QSARawKVSharding( + local_capacity=size + page_size, + world_size=cache_sharding_runtime.size, + rank=cache_sharding_runtime.rank, + ) + self.qsa_compressed_sharding = QSACompressedBlockSharding( + global_blocks=self.qsa_global_compressed_capacity, + compressed_page_size=self.qsa_compressed_page_size, + world_size=cache_sharding_runtime.size, + rank=cache_sharding_runtime.rank, + ) + self.qsa_compressed_capacity = self.qsa_compressed_sharding.local_blocks # These buffers participate in Mooncake PD transfer just like the base # KV and Mamba pools. Keep their allocation in the same memory-saver # and Mooncake custom-pool regions; otherwise MNNVL cannot resolve the @@ -252,6 +523,54 @@ def get_qsa_key_state_buffer(self, layer_id: int) -> torch.Tensor: self._transfer_full_attention_id(layer_id) ] + def set_kv_buffer( + self, + layer, + loc: torch.Tensor, + cache_k: torch.Tensor, + cache_v: torch.Tensor, + k_scale: float = 1.0, + v_scale: float = 1.0, + dcp_kv_mask: torch.Tensor | None = None, + ) -> None: + sharding = self.qsa_raw_kv_sharding + if sharding is not None: + local_layer_id = self._transfer_full_attention_id(layer.layer_id) + key_buffer = self.full_kv_pool.get_key_buffer(local_layer_id) + value_buffer = self.full_kv_pool.get_value_buffer(local_layer_id) + buffer_heads = key_buffer.shape[1] + buffer_dim = key_buffer.shape[-1] + if ( + cache_k.ndim != 3 + or cache_v.shape != cache_k.shape + or value_buffer.shape != key_buffer.shape + or cache_k.shape[1] != buffer_heads + or cache_k.shape[2] != buffer_dim + ): + raise ValueError( + "raw K/V write layout must match the owner-local cache: " + f"cache={tuple(cache_k.shape)}, buffer={tuple(key_buffer.shape)}" + ) + owner_mask, loc = sharding.local_write_targets(loc) + if dcp_kv_mask is not None: + owner_mask &= dcp_kv_mask.to(device=owner_mask.device, dtype=torch.bool) + dcp_kv_mask = owner_mask + super().set_kv_buffer( + layer, + loc, + cache_k, + cache_v, + k_scale, + v_scale, + dcp_kv_mask=dcp_kv_mask, + ) + + def move_kv_cache(self, tgt_loc: torch.Tensor, src_loc: torch.Tensor) -> None: + sharding = self.qsa_raw_kv_sharding + if sharding is not None: + src_loc, tgt_loc = sharding.local_copy_targets(src_loc, tgt_loc) + self.full_kv_pool.move_kv_cache(tgt_loc, src_loc) + def set_qsa_key_state_buffer( self, layer_id: int, loc: torch.Tensor, token_k: torch.Tensor ) -> None: @@ -284,6 +603,11 @@ def set_qsa_compressed_k_buffer( self, layer_id: int, loc: torch.Tensor, compressed_k: torch.Tensor ) -> None: buffer = self.get_qsa_compressed_k_buffer(layer_id) + if self.qsa_compressed_sharding is not None: + local_blocks = self.qsa_compressed_sharding.global_to_local_blocks(loc) + source_rows = torch.nonzero(local_blocks >= 0, as_tuple=False).flatten() + loc = local_blocks.index_select(0, source_rows) + compressed_k = compressed_k.index_select(0, source_rows) buffer[loc.long()] = compressed_k.to(buffer.dtype) @staticmethod diff --git a/python/sglang/srt/mem_cache/static_kv_storage.py b/python/sglang/srt/mem_cache/static_kv_storage.py index 23aeb1750997..1f380733f82a 100644 --- a/python/sglang/srt/mem_cache/static_kv_storage.py +++ b/python/sglang/srt/mem_cache/static_kv_storage.py @@ -335,6 +335,13 @@ def take_decode( # returns a DSV4OutCacheLoc bundle; hasattr-gated so others stay unchanged. is_dsv4 = rows is not None and hasattr(allocator, "c128_attn_allocator") extra_alloc_kwargs = {} + if page_interleave_shard_size(allocator) > 1: + assert self._batch is not None + rotation_bases = [req.kv_rotation_base for req in self._batch.reqs] + assert all(base is not None for base in rotation_bases), ( + "sharded decode requires the rotation base recorded during prefill" + ) + extra_alloc_kwargs["rotation_bases"] = rotation_bases if is_dsv4: c128_num_pages = allocator.c128_num_pages_needed( (seq_lens_cpu - 1).clamp(min=0), seq_lens_cpu diff --git a/python/sglang/srt/model_executor/model_runner.py b/python/sglang/srt/model_executor/model_runner.py index 92f1f98a51ce..90590fd5142d 100644 --- a/python/sglang/srt/model_executor/model_runner.py +++ b/python/sglang/srt/model_executor/model_runner.py @@ -1566,7 +1566,8 @@ def logical_max_total_num_tokens(self): """Request-token capacity in logical tokens, not per-rank DCP rows.""" return self.req_to_token_pool.schedulable_token_capacity( self.kv_cache_configurator.logical_token_capacity( - max_total_num_tokens=self.max_total_num_tokens + max_total_num_tokens=self.max_total_num_tokens, + allocator=self.token_to_kv_pool_allocator, ) ) diff --git a/python/sglang/srt/runtime_context.py b/python/sglang/srt/runtime_context.py index c7a5a9f5a353..4662ac9af536 100644 --- a/python/sglang/srt/runtime_context.py +++ b/python/sglang/srt/runtime_context.py @@ -248,6 +248,7 @@ class SpawnRanks(msgspec.Struct, frozen=True): ("attn_dp_rank", "attn_dp_size"), ("attn_cp_rank", "attn_cp_size"), ("moe_ep_rank", "moe_ep_size"), + ("qsa_cache_sharding_rank", "qsa_cache_sharding_size"), ) # MoE-DP may alias a wider attention-CP group, so its configured width @@ -258,6 +259,7 @@ class SpawnRanks(msgspec.Struct, frozen=True): ("attn_tp_size", "attn_tp_group"), ("attn_cp_size", "attn_cp_group"), ("moe_ep_size", "moe_ep_group"), + ("qsa_cache_sharding_size", "qsa_cache_sharding_group"), ) _UNREADABLE = object() @@ -1576,9 +1578,11 @@ def _configure_context( ), ) context._publish_role = role - # Disabled DCP has rank zero even in processes without a rank bundle. + # Disabled DCP and QSA cache sharding have rank zero without a rank bundle. if not context.parallel.dcp_enabled: context.parallel.override_permanently(attn_dcp_rank=0) + if context.parallel.qsa_cache_sharding_size == 1: + context.parallel.override_permanently(qsa_cache_sharding_rank=0) # The device is assigned by the launcher; it is not a config field. context.config_bag("device")._set( "gpu_id", ranks.gpu_id if ranks is not None else None @@ -1605,6 +1609,9 @@ def _configure_context( if parallel.dcp_enabled: placement["dcp_rank"] = placement["tp_rank"] % parallel.dcp_size placement["attn_dcp_rank"] = placement.get("dcp_rank", 0) + placement["qsa_cache_sharding_rank"] = ( + placement["tp_rank"] % parallel.qsa_cache_sharding_size + ) parallel.override_permanently(**placement) _validate_parallel(parallel, "publish") if _ROLE_NS_MODE == "record": diff --git a/test/registered/kernels/ops/attention/qsa/test_qsa.py b/test/registered/kernels/ops/attention/qsa/test_qsa.py index ccfd3ebae2ac..c96459cf22a0 100644 --- a/test/registered/kernels/ops/attention/qsa/test_qsa.py +++ b/test/registered/kernels/ops/attention/qsa/test_qsa.py @@ -36,6 +36,11 @@ QwenSparseAttnBackend, QwenSparseMultiStepDraftBackend, ) +from sglang.srt.layers.quantization.fp4_kv_cache_quant_method import ( + UnquantizedKVCacheMethod, +) +from sglang.srt.mem_cache.memory_pool import MHATokenToKVPool +from sglang.srt.mem_cache.qsa_kv_pool import QSATokenToKVPool from sglang.srt.model_executor.forward_mode import ForwardMode from sglang.test.ci.ci_register import register_cuda_ci from sglang.test.mem_cache_utils import routed @@ -326,7 +331,7 @@ def test_qsa_glue_builds_compressed_indexer(monkeypatch): class _FakeIndexer: def __init__( - self, config, layer_id, quant_config=None, prefix="", rotary_emb=None + self, config, layer_id, quant_config=None, prefix="", rotary_emb=None, **_ ): recorded.update( config=config, @@ -537,6 +542,7 @@ def test_qsa_indexer_ignores_dp_attention_token_padding(): def run(num_rows): mapping = torch.zeros(15, dtype=torch.int32) metadata = SimpleNamespace( + is_cuda_graph=False, token_to_kv_pool=None, compress_member_rows=None, decode_logical_positions=None, @@ -591,6 +597,52 @@ def run(num_rows): assert calls == {"token_rows": 15, "logical_rows": 15, "rope_rows": 15} +def test_sharded_qsa_store_uses_pool_mapping(): + calls = [] + backend = QwenSparseAttnBackend.__new__(QwenSparseAttnBackend) + backend.token_to_kv_pool = SimpleNamespace( + qsa_raw_kv_sharding=object(), + set_kv_buffer=lambda *args: calls.append(args), + ) + args = (SimpleNamespace(layer_id=7), object(), object(), object()) + backend._store_kv_cache(*args) + assert calls == [args] + + +def test_sharded_qsa_pool_is_not_fused_eligible(): + storage = MHATokenToKVPool.__new__(MHATokenToKVPool) + storage.quant_method = UnquantizedKVCacheMethod() + storage.dtype = torch.bfloat16 + pool = QSATokenToKVPool.__new__(QSATokenToKVPool) + pool.full_kv_pool = storage + pool.cache_sharding_runtime = SimpleNamespace(enabled=True) + + assert not QwenSparseAttnBackend._supports_fused_kv_pool(pool) + + +def test_distributed_decode_scores_local_candidates(monkeypatch): + if not torch.cuda.is_available(): + pytest.skip("CUDA is required") + + def fake_decode(*_, min_scored_len=0, **__): + assert min_scored_len == 0 + raise StopIteration + + monkeypatch.setattr(qsa_indexer_module, "qsa_mqa_decode", fake_decode) + indexer = SimpleNamespace(block_topk=512, distributed_topk_group=object()) + with pytest.raises(StopIteration): + QSAIndexer.select_decode_tokens( + indexer, + torch.empty(1, 1, 1, device="cuda"), + None, + None, + torch.ones(1, dtype=torch.int32, device="cuda"), + 1, + None, + None, + ) + + def test_qsa_cuda_extend_ignores_dp_attention_padding(monkeypatch): if not torch.cuda.is_available(): return @@ -640,6 +692,7 @@ def run(num_rows): def _make_paged_extend_backend(): metadata = SimpleNamespace( + is_cuda_graph=False, token_to_batch_idx=torch.zeros(3, dtype=torch.int32), sequence_lengths=torch.tensor([8], dtype=torch.int32), token_slot_table=torch.arange(16, dtype=torch.int32).reshape(1, 16), @@ -862,6 +915,7 @@ def test_qsa_speculative_pseudo_extend_is_rejected(): class _FakeQSAPool: + cache_sharding_runtime = SimpleNamespace(size=1, rank=0, group=None, enabled=False) qsa_index_kv_heads = 1 qsa_index_head_dim = 128 qsa_compressed_page_size = 64 @@ -984,7 +1038,7 @@ def select_prefill_tokens(self, *args, **kwargs): self.selected = "prefill" return torch.tensor([1]) - def select_decode_tokens(self, *args, defer_expansion=False): + def select_decode_tokens(self, *args, defer_expansion=False, **kwargs): self.selected = "decode" return torch.tensor([2]) @@ -995,6 +1049,7 @@ def update_key_state_and_compress( class _DispatchMetadata: + is_cuda_graph = False token_to_kv_pool = None out_cache_loc = None compress_member_rows = None @@ -1359,6 +1414,7 @@ class _Pool: qsa_compress_ratio = ratio # compressed slots per full-KV page (page 64 tokens / ratio 4) qsa_compressed_page_size = full_page // ratio + qsa_compressed_sharding = None qsa_block_topk = 512 def run_case(mode, bs, num_rows, seq_lens_list, extend_len, extend_lens=None): @@ -1398,9 +1454,15 @@ def make_metadata(): graph_write_locs=torch.zeros( num_rows, dtype=torch.int32, device=device ), + write_owner_mask=torch.zeros(num_rows, dtype=torch.bool, device=device), graph_compressed_page_table=torch.zeros( (num_rows, max_pages), dtype=torch.int32, device=device ), + decode_block_positions=torch.zeros( + (num_rows, max_pages * (full_page // ratio)), + dtype=torch.int32, + device=device, + ), graph_compressed_lengths=torch.zeros( num_rows, dtype=torch.int32, device=device ), @@ -1456,7 +1518,6 @@ def make_metadata(): ) for field in ( "graph_write_locs", - "graph_compressed_page_table", "graph_compressed_lengths", "decode_logical_positions", "pending_ring_slots", @@ -1465,6 +1526,14 @@ def make_metadata(): kernel_buf = getattr(kernel_metadata.indexer_metadata, field) host_buf = getattr(host_metadata.indexer_metadata, field) assert torch.equal(kernel_buf, host_buf), (mode, field) + widths = (kernel_metadata.indexer_metadata.graph_compressed_lengths + 15) // 16 + for row, width in enumerate(widths.tolist()): + assert torch.equal( + kernel_metadata.indexer_metadata.graph_compressed_page_table[ + row, :width + ], + host_metadata.indexer_metadata.graph_compressed_page_table[row, :width], + ) # Decode: lengths straddling boundaries (256 is a boundary, others not). run_case(mode=0, bs=4, num_rows=4, seq_lens_list=[255, 256, 257, 512], extend_len=0) @@ -1526,6 +1595,7 @@ def test_qsa_graph_layout_covers_speculative_rows_and_padded_tail(seq_offset): class _Pool: qsa_compress_ratio = ratio qsa_compressed_page_size = full_page // ratio + qsa_compressed_sharding = None qsa_block_topk = 512 def run_case(mode, bs, num_rows, seq_lens, extend_lens, extend_len, num_padding): @@ -1551,9 +1621,15 @@ def run_case(mode, bs, num_rows, seq_lens, extend_lens, extend_len, num_padding) req_pool_indices=torch.zeros(num_rows, dtype=torch.int32, device=device), is_cuda_graph=True, graph_write_locs=torch.zeros(num_rows, dtype=torch.int32, device=device), + write_owner_mask=torch.zeros(num_rows, dtype=torch.bool, device=device), graph_compressed_page_table=torch.zeros( (num_rows, max_pages), dtype=torch.int32, device=device ), + decode_block_positions=torch.zeros( + (num_rows, max_pages * (full_page // ratio)), + dtype=torch.int32, + device=device, + ), graph_compressed_lengths=torch.zeros( num_rows, dtype=torch.int32, device=device ), @@ -1685,9 +1761,15 @@ def metadata(): indexer = SimpleNamespace( graph_compressed_lengths=torch.empty(bs, dtype=torch.int32, device="cuda"), graph_write_locs=torch.empty(bs, dtype=torch.int32, device="cuda"), - graph_compressed_page_table=torch.empty( + write_owner_mask=torch.empty(bs, dtype=torch.bool, device="cuda"), + graph_compressed_page_table=torch.zeros( (bs, pages), dtype=torch.int32, device="cuda" ), + decode_block_positions=torch.zeros( + (bs, pages * (full_page // ratio)), + dtype=torch.int32, + device="cuda", + ), decode_logical_positions=torch.empty(bs, dtype=torch.int32, device="cuda"), pending_ring_slots=torch.empty(bs, dtype=torch.int64, device="cuda"), graph_ring_group_locs=torch.empty( diff --git a/test/registered/kernels/ops/attention/qsa/test_qsa_cache_sharding.py b/test/registered/kernels/ops/attention/qsa/test_qsa_cache_sharding.py new file mode 100644 index 000000000000..5a5256f95293 --- /dev/null +++ b/test/registered/kernels/ops/attention/qsa/test_qsa_cache_sharding.py @@ -0,0 +1,154 @@ +import os + +import pytest +import torch +import torch.distributed as dist +import torch.multiprocessing as mp + +import sglang.srt.layers.attention.qsa.distributed_topk as distributed_topk +from sglang.srt.layers.attention.qsa.cache_sharding import QSACacheShardingStaticBuffers +from sglang.srt.layers.attention.qsa.owner_sparse_attn import owner_sparse_attention +from sglang.srt.mem_cache.qsa_kv_pool import QSARawKVSharding +from sglang.test.ci.ci_register import register_cuda_ci +from sglang.test.test_utils import find_available_port + +register_cuda_ci(est_time=30, stage="base-b-kernel-unit", runner_config="4-gpu-b200") + + +class _Group: + def __init__(self, rank, world_size): + self.rank_in_group, self.world_size = rank, world_size + self.device_group = dist.group.WORLD + + def all_to_all_single(self, output, value): + dist.all_to_all_single(output, value) + + def all_gather(self, value, dim=0): + values = [torch.empty_like(value) for _ in range(self.world_size)] + dist.all_gather(values, value) + return torch.cat(values, dim=dim) + + +def _capture(run): + stream = torch.cuda.Stream() + stream.wait_stream(torch.cuda.current_stream()) + with torch.cuda.stream(stream): + run() + torch.cuda.current_stream().wait_stream(stream) + torch.cuda.synchronize() + graph = torch.cuda.CUDAGraph() + with torch.cuda.graph(graph): + outputs = run() + return graph, outputs + + +def _worker(rank, world_size, port): + os.environ.update( + MASTER_ADDR="127.0.0.1", + MASTER_PORT=str(port), + SGLANG_QSA_CANDIDATE_SYMM_MEM="1", + ) + torch.cuda.set_device(rank) + dist.init_process_group("nccl", rank=rank, world_size=world_size) + try: + device = torch.device("cuda", rank) + group = _Group(rank, world_size) + rows, topk, head_dim = 2, 64, 128 + torch.manual_seed(43527) + + def randn(*shape): + return torch.randn(*shape, dtype=torch.bfloat16, device=device) + + def expand(values): + return values.expand(rows, -1).contiguous() + + global_tokens = topk * world_size + logical = expand( + torch.arange( + rank, global_tokens, world_size, dtype=torch.int32, device=device + ) + ) + logits = randn(rows, topk).float() + local_ids = expand(torch.arange(topk, dtype=torch.int32, device=device)) + row_starts = torch.zeros(rows, dtype=torch.int32, device=device) + full_q = randn(rows, world_size, head_dim) + local_q = full_q[:, rank : rank + 1].contiguous() + full_k = randn(global_tokens, 1, head_dim) + full_v = randn(global_tokens, 1, head_dim) + local_k = full_k[rank::world_size].contiguous() + local_v = full_v[rank::world_size].contiguous() + buffers = QSACacheShardingStaticBuffers( + max_rows=rows, + topk=topk, + num_heads=world_size, + head_dim=head_dim, + world_size=world_size, + device=device, + dtype=torch.bfloat16, + candidate_group=group, + ) + local_shard = QSARawKVSharding(topk, world_size, rank) + full_shard = QSARawKVSharding(global_tokens, 1, 0) + + def publish(): + return distributed_topk.fused_publish_and_merge_qsa_topk_candidates( + logits, + local_ids, + row_starts, + logical, + group=group, + topk=topk, + static_buffers=buffers, + ) + + def run(): + scores, slots = publish() + output, lse = owner_sparse_attention( + local_q, + local_k, + local_v, + slots, + local_shard, + group=group, + static_buffers=buffers, + ) + return scores, slots, output, lse + + def reference(): + scores, slots = distributed_topk.gather_and_merge_qsa_topk_candidates( + logits, logical, group=group, topk=topk + ) + output, lse = owner_sparse_attention( + full_q, full_k, full_v, slots, full_shard + ) + return scores, slots, output[:, rank : rank + 1], lse[:, rank : rank + 1] + + def check(actual, expected): + torch.testing.assert_close(actual[0], expected[0], rtol=0, atol=0) + torch.testing.assert_close(actual[1], expected[1], rtol=0, atol=0) + torch.testing.assert_close(actual[2], expected[2], rtol=2e-2, atol=2e-2) + torch.testing.assert_close(actual[3], expected[3], rtol=2e-3, atol=2e-3) + + check(run(), reference()) + dist.barrier(device_ids=[rank]) + graph, outputs = _capture(publish) + logits.add_(rank + 0.25) + graph.replay() + torch.cuda.synchronize() + expected = distributed_topk.gather_and_merge_qsa_topk_candidates( + logits, logical, group=group, topk=topk + ) + for actual, reference in zip(outputs, expected): + torch.testing.assert_close(actual, reference, rtol=0, atol=0) + finally: + dist.destroy_process_group() + + +def test_qsa_multimem_owner_pipeline_eager_and_cuda_graph(): + if torch.cuda.device_count() < 4 or torch.cuda.get_device_capability()[0] != 9: + pytest.skip("four SM90 CUDA devices are required") + mp.spawn(_worker, args=(4, find_available_port(29500)), nprocs=4, join=True) + + +if __name__ == "__main__": + raise SystemExit(pytest.main([__file__, "-v"])) diff --git a/test/registered/unit/managers/test_dcp_logical_capacity.py b/test/registered/unit/managers/test_dcp_logical_capacity.py index 9578fe69f03d..e61caa4b0a5e 100644 --- a/test/registered/unit/managers/test_dcp_logical_capacity.py +++ b/test/registered/unit/managers/test_dcp_logical_capacity.py @@ -120,11 +120,12 @@ def test_sharded_worker_request_limits(self): runner.max_total_num_tokens = 64 allocator = PageInterleavePoolAllocator.__new__(PageInterleavePoolAllocator) allocator.shard_size = 3 + allocator.size = 64 * 3 runner.token_to_kv_pool_allocator = allocator info = TpModelWorker.get_worker_info(worker) - self.assertEqual(info[0], 64) + self.assertEqual(info[0], 64 * 3) self.assertEqual(info[4], 64 * 3 - 1) self.assertEqual(info[5], 64 * 3 - 6) diff --git a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py index 22da06a9d2f5..2f46586b0d2b 100644 --- a/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py +++ b/test/registered/unit/managers/test_scheduler_init_req_max_new_tokens.py @@ -21,7 +21,7 @@ class TestSchedulerInitReqMaxNewTokens(unittest.TestCase): 1. context: input_len + max_new_tokens < max_req_len 2. admission budget (PrefillAdder): ceil_page(input_len) + max_new_tokens + page_size * shard_widening - < max_total_num_tokens * shard_widening + < max_total_num_tokens 3. env limit: <= SGLANG_MAX_NEW_TOKENS_LIMIT when set and positive 4. never above the requested value 5. min_new_tokens <= max_new_tokens afterwards @@ -100,7 +100,7 @@ def satisfies_rules(candidate: int) -> bool: context_ok = input_len + candidate < scheduler.max_req_len budget_ok = ( paged_input_len + candidate + page_size * shard_widening - < scheduler.max_total_num_tokens * shard_widening + < scheduler.max_total_num_tokens ) limit_ok = not limit_active or candidate <= limit requested_ok = requested is None or candidate <= requested @@ -169,7 +169,7 @@ def test_dcp_budget_is_already_in_logical_tokens(self): self.assertEqual(self._init_and_check(scheduler, req), 11) def test_sharded_budget_reserves_one_page_per_shard(self): - max_total_num_tokens, page_size, input_len, shard_widening = 8, 4, 8, 3 + max_total_num_tokens, page_size, input_len, shard_widening = 24, 4, 8, 3 scheduler = self._new_scheduler( max_req_len=128, max_total_num_tokens=max_total_num_tokens, @@ -180,15 +180,14 @@ def test_sharded_budget_reserves_one_page_per_shard(self): max_new_tokens = self._init_and_check(scheduler, req) - widened_capacity = max_total_num_tokens * shard_widening per_req_overhead = page_size * shard_widening self.assertEqual( max_new_tokens, - widened_capacity - input_len - per_req_overhead - 1, + max_total_num_tokens - input_len - per_req_overhead - 1, ) self.assertLess( input_len + max_new_tokens + per_req_overhead, - widened_capacity, + max_total_num_tokens, ) def test_min_new_tokens_clamped_to_limit(self): diff --git a/test/registered/unit/mem_cache/test_page_interleave_shard.py b/test/registered/unit/mem_cache/test_page_interleave_shard.py index b4b1e272589e..49df2092c88e 100644 --- a/test/registered/unit/mem_cache/test_page_interleave_shard.py +++ b/test/registered/unit/mem_cache/test_page_interleave_shard.py @@ -916,10 +916,6 @@ def test_unsupported_paths_fail_loud(self): alloc = _make_allocator() with self.assertRaises(NotImplementedError): alloc.alloc(GS) - with self.assertRaises(NotImplementedError): - alloc.alloc_decode( - torch.tensor([PS + 1]), torch.tensor([PS + 1]), torch.tensor([PS - 1]) - ) def test_batch_alloc_per_request_rotation(self): """bs > 1: each request draws its own cyclic run; out_cache_loc is diff --git a/test/registered/unit/test_runtime_context.py b/test/registered/unit/test_runtime_context.py index 4ee993244f1d..91c12d2dd8b4 100644 --- a/test/registered/unit/test_runtime_context.py +++ b/test/registered/unit/test_runtime_context.py @@ -107,6 +107,7 @@ def _parallel_state(): GROUP_STAMPS = { "tp_group": "_TP", "dcp_group": "_DCP", + "qsa_cache_sharding_group": "_QSA_CACHE_SHARDING", "pp_group": "_PP", "moe_ep_group": "_MOE_EP", "moe_dp_group": "_MOE_DP",