From 97a0dd358aa79d0f6100a6c5cb5ac47e792edfe9 Mon Sep 17 00:00:00 2001 From: MaxwellF1 Date: Mon, 31 Aug 2026 19:12:27 -0700 Subject: [PATCH] Add: support packed multi-request DSpark prefill - Add packed query boundaries and rank-local request ownership to the DSpark DSA-CP prefill ABI. - Select request-scoped cache, compressor-state, indexer, and sparse attention rows across SWA, HCA, and CSA. - Preserve TP-aligned padding and empty-rank behavior while serializing HCA streaming work by request. - Keep indexer state publication ordered across physical tiles without treating fence values as row-validity predicates. - Thread packed-request metadata through the single-layer and 43-layer forward paths, with fixtures that cross TP-rank request boundaries. - Validate request-indexed fixture axes and reject incompatible fixed ragged-fixture overrides before dispatch. --- .../prefill_compressor_ratio128.py | 76 +-- .../prefill_compressor_ratio4.py | 74 +-- .../deepseek_v4_flash_dspark/prefill_csa.py | 476 +++++++++++++----- .../deepseek_v4_flash_dspark/prefill_fwd.py | 189 +++++-- .../deepseek_v4_flash_dspark/prefill_hca.py | 441 ++++++++++++---- .../prefill_indexer.py | 148 ++++-- .../prefill_indexer_compressor.py | 147 +++--- .../deepseek_v4_flash_dspark/prefill_layer.py | 79 ++- .../prefill_metadata.py | 116 +++++ .../prefill_sparse_attn.py | 103 ++-- .../deepseek_v4_flash_dspark/prefill_swa.py | 242 ++++++--- 11 files changed, 1533 insertions(+), 558 deletions(-) create mode 100644 models/deepseek_v4_flash_dspark/prefill_metadata.py diff --git a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py index 0f24b7b71..f46d4d62d 100644 --- a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py +++ b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio128.py @@ -6,7 +6,7 @@ # INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. # See LICENSE in the root of the software repository for the full text of the License. # ----------------------------------------------------------------------------------------------------------- -"""DeepSeek-V4 token-major prefill compressor, ratio=128, single request of <=T tokens.""" +"""DeepSeek-V4 packed prefill compressor for the ratio-128 state cache.""" import pypto.language as pl @@ -17,6 +17,7 @@ HCA_STATE_PHYSICAL_BLOCKS, PREFILL_SEQ, ) +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN @@ -394,10 +395,11 @@ def _prefill_compressor_ratio128_tile( @pl.jit.inline(auto_scope=False) def prefill_compressor_ratio128( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], wkv: pl.Tensor[[OUT_DIM, D], pl.BF16], wgate: pl.Tensor[[OUT_DIM, D], pl.BF16], ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32], @@ -409,8 +411,8 @@ def prefill_compressor_ratio128( cmp_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], ): - """Compress one physical dynamic token run through ordered 512-row state tiles.""" - t_dim = pl.tensor.dim(x, 0) + """Compress packed requests independently through ordered 512-row state tiles.""" + request_count = pl.tensor.dim(query_start_loc, 0) - 1 rope_dup_idx_template = pl.create_tensor([HCA_C128_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.INT32) rope_swap_idx_template = pl.create_tensor([HCA_C128_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.INT32) rope_sign_template = pl.create_tensor([HCA_C128_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.FP32) @@ -443,40 +445,46 @@ def prefill_compressor_ratio128( state_order_fence = pl.create_tensor([1], dtype=pl.INT32) with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_hca_c128_state_order_init"): pl.write(state_order_fence, [0], pl.cast(0, pl.INT32)) - for tile_base in pl.range(0, t_dim, PREFILL_STATE_TILE): - tile_rows = pl.min(PREFILL_STATE_TILE, t_dim - tile_base) - with pl.scope(): - _prefill_compressor_ratio128_tile( - x, - compress_state, - compress_state_block_table, - wkv, - wgate, - ape, - norm_w, - cmp_freqs_cos, - cmp_freqs_sin, - cmp_kv, - position_ids, - cmp_slot_mapping, - state_slot_mapping, - rope_dup_idx_template, - rope_swap_idx_template, - rope_sign_template, - state_order_fence, - tile_base, - tile_rows, - ) + for request in pl.range(request_count): + request_start = pl.cast(pl.read(query_start_loc, [request]), pl.INDEX) + request_end = pl.cast(pl.read(query_start_loc, [request + 1]), pl.INDEX) + request_table = compress_state_block_table[request] + for request_offset in pl.range(0, request_end - request_start, PREFILL_STATE_TILE): + tile_base = request_start + request_offset + tile_rows = pl.min(PREFILL_STATE_TILE, request_end - tile_base) + with pl.scope(): + _prefill_compressor_ratio128_tile( + x, + compress_state, + request_table, + wkv, + wgate, + ape, + norm_w, + cmp_freqs_cos, + cmp_freqs_sin, + cmp_kv, + position_ids, + cmp_slot_mapping, + state_slot_mapping, + rope_dup_idx_template, + rope_swap_idx_template, + rope_sign_template, + state_order_fence, + tile_base, + tile_rows, + ) return cmp_kv, compress_state @pl.jit def prefill_compressor_ratio128_test( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], wkv: pl.Tensor[[OUT_DIM, D], pl.BF16], wgate: pl.Tensor[[OUT_DIM, D], pl.BF16], ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32], @@ -489,7 +497,9 @@ def prefill_compressor_ratio128_test( state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], ): x.bind_dynamic(0, T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) compress_state.bind_dynamic(0, STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) cmp_kv.bind_dynamic(0, CMP_BLOCK_NUM_DYN) cmp_freqs_cos.bind_dynamic(0, T_DYN) cmp_freqs_sin.bind_dynamic(0, T_DYN) @@ -499,6 +509,7 @@ def prefill_compressor_ratio128_test( return prefill_compressor_ratio128( x, + query_start_loc, compress_state, compress_state_block_table, wkv, @@ -526,7 +537,7 @@ def golden_prefill_compressor_ratio128(tensors): ) kv_state_flat = compress_state_flat[:, :OUT_DIM] score_state_flat = compress_state_flat[:, OUT_DIM:] - state_block_table = tensors["compress_state_block_table"] + state_block_table = tensors["compress_state_block_table"][0] cmp_kv_flat = tensors["cmp_kv"].view(CMP_MAX_BLOCKS * BLOCK_SIZE, HEAD_DIM) def state_row(abs_pos): @@ -595,7 +606,7 @@ def build_tensor_specs(start_pos: int = START_POS, token_count: int = PREFILL_SE def init_compress_state_block_table(): logical_blocks = torch.arange(HCA_STATE_MAX_BLOCKS, dtype=torch.int64) - return ((logical_blocks * 17 + 3) % HCA_STATE_PHYSICAL_BLOCKS).to(torch.int32) + return ((logical_blocks * 17 + 3) % HCA_STATE_PHYSICAL_BLOCKS).to(torch.int32).unsqueeze(0) def state_row(abs_pos): if abs_pos < 0 or abs_pos >= MAX_SEQ_LEN: @@ -671,6 +682,7 @@ def init_state_slot_mapping(): return [ TensorSpec("x", [token_count, D], torch.bfloat16, init_value=init_x), + TensorSpec("query_start_loc", [2], torch.int32, init_value=torch.tensor([0, token_count], dtype=torch.int32)), TensorSpec( "compress_state", [HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], @@ -680,7 +692,7 @@ def init_state_slot_mapping(): ), TensorSpec( "compress_state_block_table", - [HCA_STATE_MAX_BLOCKS], + [1, HCA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table, ), diff --git a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py index ae791ef02..c753e0396 100644 --- a/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py +++ b/models/deepseek_v4_flash_dspark/prefill_compressor_ratio4.py @@ -18,6 +18,7 @@ FP32_NEG_INF, PREFILL_SEQ, ) +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN @@ -483,10 +484,11 @@ def _prefill_compressor_ratio4_tile( @pl.jit.inline(auto_scope=False) def compressor_ratio4( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], wkv: pl.Tensor[[OUT_DIM, D], pl.BF16], wgate: pl.Tensor[[OUT_DIM, D], pl.BF16], ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32], @@ -498,8 +500,8 @@ def compressor_ratio4( cmp_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], ): - """Compress one physical dynamic token run through ordered 512-row state tiles.""" - t_dim = pl.tensor.dim(x, 0) + """Compress packed requests independently through ordered 512-row state tiles.""" + request_count = pl.tensor.dim(query_start_loc, 0) - 1 rope_dup_idx_template = pl.create_tensor([PACKED_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.INT32) rope_swap_idx_template = pl.create_tensor([PACKED_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.INT32) rope_sign_template = pl.create_tensor([PACKED_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.FP32) @@ -526,30 +528,35 @@ def compressor_ratio4( state_order_fence = pl.create_tensor([1], dtype=pl.INT32) with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_c4_state_order_init"): pl.write(state_order_fence, [0], pl.cast(0, pl.INT32)) - for tile_base in pl.range(0, t_dim, PREFILL_STATE_TILE): - tile_rows = pl.min(PREFILL_STATE_TILE, t_dim - tile_base) - with pl.scope(): - _prefill_compressor_ratio4_tile( - x, - compress_state, - compress_state_block_table, - wkv, - wgate, - ape, - norm_w, - cmp_freqs_cos, - cmp_freqs_sin, - cmp_kv, - position_ids, - cmp_slot_mapping, - state_slot_mapping, - rope_dup_idx_template, - rope_swap_idx_template, - rope_sign_template, - state_order_fence, - tile_base, - tile_rows, - ) + for request in pl.range(request_count): + request_start = pl.cast(pl.read(query_start_loc, [request]), pl.INDEX) + request_end = pl.cast(pl.read(query_start_loc, [request + 1]), pl.INDEX) + request_table = compress_state_block_table[request] + for request_offset in pl.range(0, request_end - request_start, PREFILL_STATE_TILE): + tile_base = request_start + request_offset + tile_rows = pl.min(PREFILL_STATE_TILE, request_end - tile_base) + with pl.scope(): + _prefill_compressor_ratio4_tile( + x, + compress_state, + request_table, + wkv, + wgate, + ape, + norm_w, + cmp_freqs_cos, + cmp_freqs_sin, + cmp_kv, + position_ids, + cmp_slot_mapping, + state_slot_mapping, + rope_dup_idx_template, + rope_swap_idx_template, + rope_sign_template, + state_order_fence, + tile_base, + tile_rows, + ) return cmp_kv, compress_state @@ -562,7 +569,7 @@ def golden_prefill_compressor_ratio4(tensors): compress_state_flat = tensors["compress_state"].view(-1, COMPRESS_STATE_DIM) kv_state_flat = compress_state_flat[:, :OUT_DIM] score_state_flat = compress_state_flat[:, OUT_DIM:] - state_block_table = tensors["compress_state_block_table"] + state_block_table = tensors["compress_state_block_table"][0] wkv = tensors["wkv"].float() wgate = tensors["wgate"].float() ape = tensors["ape"] @@ -661,10 +668,11 @@ def state_row(abs_pos): @pl.jit def prefill_compressor_ratio4_test( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], wkv: pl.Tensor[[OUT_DIM, D], pl.BF16], wgate: pl.Tensor[[OUT_DIM, D], pl.BF16], ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32], @@ -677,7 +685,9 @@ def prefill_compressor_ratio4_test( state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], ): x.bind_dynamic(0, T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) compress_state.bind_dynamic(0, STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) cmp_kv.bind_dynamic(0, CMP_BLOCK_NUM_DYN) cmp_freqs_cos.bind_dynamic(0, T_DYN) cmp_freqs_sin.bind_dynamic(0, T_DYN) @@ -687,6 +697,7 @@ def prefill_compressor_ratio4_test( return compressor_ratio4( x, + query_start_loc, compress_state, compress_state_block_table, wkv, @@ -716,7 +727,7 @@ def build_tensor_specs(start_pos: int = START_POS, token_count: int = PREFILL_SE def init_compress_state_block_table(): logical_blocks = torch.arange(CSA_STATE_MAX_BLOCKS, dtype=torch.int64) - return ((logical_blocks * 17 + 3) % CSA_STATE_PHYSICAL_BLOCKS).to(torch.int32) + return ((logical_blocks * 17 + 3) % CSA_STATE_PHYSICAL_BLOCKS).to(torch.int32).unsqueeze(0) def state_row(abs_pos): if abs_pos < 0 or abs_pos >= MAX_SEQ_LEN: @@ -796,6 +807,7 @@ def init_state_slot_mapping(): return [ TensorSpec("x", [token_count, D], torch.bfloat16, init_value=init_x), + TensorSpec("query_start_loc", [2], torch.int32, init_value=torch.tensor([0, token_count], dtype=torch.int32)), TensorSpec( "compress_state", [CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], @@ -805,7 +817,7 @@ def init_state_slot_mapping(): ), TensorSpec( "compress_state_block_table", - [CSA_STATE_MAX_BLOCKS], + [1, CSA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table, ), diff --git a/models/deepseek_v4_flash_dspark/prefill_csa.py b/models/deepseek_v4_flash_dspark/prefill_csa.py index 3c367dc7e..540f2bed0 100644 --- a/models/deepseek_v4_flash_dspark/prefill_csa.py +++ b/models/deepseek_v4_flash_dspark/prefill_csa.py @@ -49,6 +49,7 @@ INNER_STATE_MAX_BLOCKS, prefill_indexer_compressor, ) +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN from qkv_proj_rope import golden_qkv_proj_rope, qkv_proj_rope from rmsnorm import golden_rms_norm, rms_norm from prefill_sparse_attn import ( @@ -148,6 +149,7 @@ @pl.jit.inline def prefill_attention_csa( x_hc: pl.Tensor[[T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -169,7 +171,7 @@ def prefill_attention_csa( compress_state: pl.Tensor[ [MAIN_STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32 ], - compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], hadamard_idx: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], idx_wq_b_scale: pl.Tensor[[IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -181,16 +183,17 @@ def prefill_attention_csa( inner_compress_state: pl.Tensor[ [INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32 ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8], idx_kv_scale: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids: pl.Tensor[[T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], idx_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], @@ -201,7 +204,7 @@ def prefill_attention_csa( wo_b_scale: pl.Tensor[[D], pl.FP32], x_out: pl.Tensor[[T_DYN, HC_MULT, D], pl.FP32], ): - """Run CSA for one contiguous request with increasing physical position rows.""" + """Run CSA over one packed ragged prefill stream.""" t_dim = pl.tensor.dim(x_hc, 0) x_mixed = pl.create_tensor([t_dim, D], dtype=pl.BF16) post = pl.create_tensor([t_dim, HC_MULT], dtype=pl.FP32) @@ -246,6 +249,7 @@ def prefill_attention_csa( compressor_ratio4( x_normed, + query_start_loc, compress_state, compress_state_block_table, cmp_wkv, @@ -274,6 +278,7 @@ def prefill_attention_csa( cmp_topk_indices = pl.create_tensor([t_dim, IDX_TOPK], dtype=pl.INT32) idx_kv_cache_out, idx_kv_scale_out, cmp_topk_indices = prefill_indexer( x_normed, + query_start_loc, qr, qr_scale, idx_wq_b, @@ -295,6 +300,7 @@ def prefill_attention_csa( idx_block_table, cmp_topk_indices, position_ids, + local_request_ids, idx_slot_mapping, inner_state_slot_mapping, ) @@ -309,29 +315,32 @@ def prefill_attention_csa( swa_row = pl.full([1, WIN], dtype=pl.INT32, value=-1) mask_row = pl.full([1, VALID_BLOCK_MASK_COLS], dtype=pl.INT32, value=0) if t_idx < t_dim: - abs_pos = pl.read(position_ids, [t_idx]) - # Sparse-block liveness from the dense TopK prefix. - visible_cmp = pl.min((abs_pos + 1) // COMPRESS_RATIO, pl.cast(IDX_TOPK, pl.INT32)) - for mask_sb in pl.unroll(PREFILL_ATTN_BLOCKS): - cmp_lo = pl.max(mask_sb * PREFILL_ATTN_TILE - WIN, pl.cast(0, pl.INT32)) - cmp_hi = pl.min( - (mask_sb + 1) * PREFILL_ATTN_TILE - WIN, pl.cast(SPARSE_CMP_BIAS_COLS, pl.INT32) - ) - if cmp_lo < cmp_hi: - if visible_cmp > cmp_lo: - pl.write(mask_row, [0, mask_sb], pl.cast(1, pl.INT32)) - window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) - key_start_abs = abs_pos + 1 - window_valid - for win_col in pl.range(WIN): - win_col_i32 = pl.cast(win_col, pl.INT32) - if win_col_i32 < window_valid: - key_abs = key_start_abs + win_col_i32 - blk_slot = key_abs // BLOCK_SIZE - blk = pl.read(ori_block_table, [pl.cast(blk_slot, pl.INDEX)]) - if blk >= 0: - row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) - pl.write(swa_row, [0, win_col], row) - pl.write(mask_row, [0, win_col // PREFILL_ATTN_TILE], pl.cast(1, pl.INT32)) + request_id = pl.read(local_request_ids, [t_idx]) + if request_id >= 0: + abs_pos = pl.read(position_ids, [t_idx]) + # Sparse-block liveness from the dense TopK prefix. + visible_cmp = pl.min((abs_pos + 1) // COMPRESS_RATIO, pl.cast(IDX_TOPK, pl.INT32)) + for mask_sb in pl.unroll(PREFILL_ATTN_BLOCKS): + cmp_lo = pl.max(mask_sb * PREFILL_ATTN_TILE - WIN, pl.cast(0, pl.INT32)) + cmp_hi = pl.min( + (mask_sb + 1) * PREFILL_ATTN_TILE - WIN, + pl.cast(SPARSE_CMP_BIAS_COLS, pl.INT32), + ) + if cmp_lo < cmp_hi: + if visible_cmp > cmp_lo: + pl.write(mask_row, [0, mask_sb], pl.cast(1, pl.INT32)) + window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) + key_start_abs = abs_pos + 1 - window_valid + for win_col in pl.range(WIN): + win_col_i32 = pl.cast(win_col, pl.INT32) + if win_col_i32 < window_valid: + key_abs = key_start_abs + win_col_i32 + blk_slot = key_abs // BLOCK_SIZE + blk = pl.read(ori_block_table, [request_id, pl.cast(blk_slot, pl.INDEX)]) + if blk >= 0: + row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) + pl.write(swa_row, [0, win_col], row) + pl.write(mask_row, [0, win_col // PREFILL_ATTN_TILE], pl.cast(1, pl.INT32)) swa_indices[t_idx : t_idx + 1, 0:WIN] = swa_row valid_block_mask[t_idx : t_idx + 1, 0:VALID_BLOCK_MASK_COLS] = mask_row @@ -340,7 +349,7 @@ def prefill_attention_csa( attn_out = sparse_attn_physical( q, kv_cache, swa_indices, - cmp_kv, cmp_block_table, cmp_topk_indices, + cmp_kv, cmp_block_table, local_request_ids, cmp_topk_indices, valid_block_mask, attn_sink, freqs_cos, freqs_sin, wo_a, wo_b, wo_b_scale, @@ -363,6 +372,7 @@ def prefill_attention_csa( @pl.jit def prefill_attention_csa_test( x_hc: pl.Tensor[[T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -384,7 +394,7 @@ def prefill_attention_csa_test( compress_state: pl.InOut[ pl.Tensor[[MAIN_STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], hadamard_idx: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], idx_wq_b_scale: pl.Tensor[[IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -396,16 +406,17 @@ def prefill_attention_csa_test( inner_compress_state: pl.InOut[ pl.Tensor[[INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32] ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], cmp_kv: pl.InOut[pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32]], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids: pl.Tensor[[T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], idx_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], @@ -417,18 +428,25 @@ def prefill_attention_csa_test( x_out: pl.Out[pl.Tensor[[T_DYN, HC_MULT, D], pl.FP32]], ): x_hc.bind_dynamic(0, T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) freqs_cos.bind_dynamic(0, T_DYN) freqs_sin.bind_dynamic(0, T_DYN) cmp_freqs_cos.bind_dynamic(0, T_DYN) cmp_freqs_sin.bind_dynamic(0, T_DYN) compress_state.bind_dynamic(0, MAIN_STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) inner_compress_state.bind_dynamic(0, INNER_STATE_BLOCK_NUM_DYN) + inner_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) kv_cache.bind_dynamic(0, ORI_BLOCK_NUM_DYN) + ori_block_table.bind_dynamic(0, REQUESTS_DYN) ori_slot_mapping.bind_dynamic(0, T_DYN) cmp_kv.bind_dynamic(0, CMP_BLOCK_NUM_DYN) + cmp_block_table.bind_dynamic(0, REQUESTS_DYN) idx_kv_cache.bind_dynamic(0, IDX_BLOCK_NUM_DYN) idx_kv_scale.bind_dynamic(0, IDX_BLOCK_NUM_DYN) + idx_block_table.bind_dynamic(0, REQUESTS_DYN) position_ids.bind_dynamic(0, T_DYN) + local_request_ids.bind_dynamic(0, T_DYN) cmp_slot_mapping.bind_dynamic(0, T_DYN) idx_slot_mapping.bind_dynamic(0, T_DYN) state_slot_mapping.bind_dynamic(0, T_DYN) @@ -437,6 +455,7 @@ def prefill_attention_csa_test( prefill_attention_csa( x_hc, + query_start_loc, hc_attn_fn, hc_attn_scale, hc_attn_base, @@ -476,6 +495,7 @@ def prefill_attention_csa_test( idx_kv_scale, idx_block_table, position_ids, + local_request_ids, cmp_slot_mapping, idx_slot_mapping, state_slot_mapping, @@ -545,28 +565,36 @@ def golden_prefill_attention_csa(tensors): } ) - golden_prefill_compressor_ratio4( - { - "x": x_normed.view(token_count, D), - "compress_state": tensors["compress_state"], - "compress_state_block_table": tensors["compress_state_block_table"], - "wkv": tensors["cmp_wkv"], - "wgate": tensors["cmp_wgate"], - "ape": tensors["cmp_ape"], - "norm_w": tensors["cmp_norm_w"], - "cmp_freqs_cos": tensors["cmp_freqs_cos"], - "cmp_freqs_sin": tensors["cmp_freqs_sin"], - "cmp_kv": tensors["cmp_kv"], - "position_ids": tensors["position_ids"], - "cmp_slot_mapping": tensors["cmp_slot_mapping"], - "state_slot_mapping": tensors["state_slot_mapping"], - } - ) + query_start_loc = tensors["query_start_loc"] + for request in range(query_start_loc.numel() - 1): + request_start = int(query_start_loc[request].item()) + request_end = int(query_start_loc[request + 1].item()) + if request_end <= request_start: + continue + request_rows = slice(request_start, request_end) + golden_prefill_compressor_ratio4( + { + "x": x_normed[request_rows].view(request_end - request_start, D), + "compress_state": tensors["compress_state"], + "compress_state_block_table": tensors["compress_state_block_table"][request : request + 1], + "wkv": tensors["cmp_wkv"], + "wgate": tensors["cmp_wgate"], + "ape": tensors["cmp_ape"], + "norm_w": tensors["cmp_norm_w"], + "cmp_freqs_cos": tensors["cmp_freqs_cos"][request_rows], + "cmp_freqs_sin": tensors["cmp_freqs_sin"][request_rows], + "cmp_kv": tensors["cmp_kv"], + "position_ids": tensors["position_ids"][request_rows], + "cmp_slot_mapping": tensors["cmp_slot_mapping"][request_rows], + "state_slot_mapping": tensors["state_slot_mapping"][request_rows], + } + ) idx_cos = rope_cos_t[:, :HALF_ROPE].float().contiguous() idx_sin = rope_sin_t[:, :HALF_ROPE].float().contiguous() cmp_topk_indices = golden_prefill_indexer_core( { "x": x_normed.view(token_count, D), + "query_start_loc": tensors["query_start_loc"], "qr": qr, "qr_scale": qr_scale, "wq_b": tensors["idx_wq_b"], @@ -587,6 +615,7 @@ def golden_prefill_attention_csa(tensors): "idx_kv_scale": tensors["idx_kv_scale"], "idx_block_table": tensors["idx_block_table"], "position_ids": tensors["position_ids"], + "local_request_ids": tensors["local_request_ids"], "idx_slot_mapping": tensors["idx_slot_mapping"], "inner_state_slot_mapping": tensors["inner_state_slot_mapping"], } @@ -603,12 +632,16 @@ def assemble_swa_indices(): swa_idx = torch.full((token_count, WIN), -1, dtype=torch.int32) pos = tensors["position_ids"] ori_table = tensors["ori_block_table"] + request_ids = tensors["local_request_ids"] for t in range(token_count): + request_id = int(request_ids[t].item()) + if request_id < 0: + continue abs_pos = int(pos[t].item()) window_valid = min(WIN, abs_pos + 1) key_start_abs = abs_pos + 1 - window_valid for k, key_abs in enumerate(range(key_start_abs, abs_pos + 1)): - row = cache_row_from_table(ori_table, key_abs) + row = cache_row_from_table(ori_table[request_id], key_abs) if row >= 0: swa_idx[t, k] = row return swa_idx @@ -629,6 +662,7 @@ def assemble_swa_indices(): "swa_indices": swa_indices, "cmp_kv": tensors["cmp_kv"], "cmp_block_table": tensors["cmp_block_table"], + "local_request_ids": tensors["local_request_ids"], "cmp_indices": cmp_indices, "attn_sink": tensors["attn_sink"], "freqs_cos": rope_cos_t, @@ -856,7 +890,7 @@ def init_cmp_norm_w(): state_table = _state_block_table(CSA_STATE_MAX_BLOCKS, CSA_STATE_PHYSICAL_BLOCKS) def init_compress_state_block_table(): - return state_table.clone() + return state_table.clone().unsqueeze(0) def init_compress_state(): state = torch.zeros(CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM) @@ -900,7 +934,7 @@ def init_inner_norm_w(): ) def init_inner_compress_state_block_table(): - return inner_state_table.clone() + return inner_state_table.clone().unsqueeze(0) def init_inner_compress_state(): state = torch.zeros(INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM) @@ -950,7 +984,7 @@ def init_kv_cache(): return cache def init_ori_block_table(): - return ori_block_table.clone() + return ori_block_table.clone().unsqueeze(0) def init_ori_slot_mapping(): return paged_rows(ori_block_table, token_pos().to(torch.int64)) @@ -963,14 +997,17 @@ def init_cmp_kv(): return cache def init_cmp_block_table(): - return cmp_block_table.clone() + return cmp_block_table.clone().unsqueeze(0) def init_idx_block_table(): - return idx_block_table.clone() + return idx_block_table.clone().unsqueeze(0) def init_position_ids(): return token_pos() + def init_local_request_ids(): + return torch.zeros(token_count, dtype=torch.int32) + def init_cmp_slot_mapping(): mapping = torch.full((token_count,), -1, dtype=torch.int64) token_ids, cmp_slots = cmp_write_records() @@ -1010,6 +1047,7 @@ def init_wo_b(): return [ TensorSpec("x_hc", [token_count, HC_MULT, D], torch.float32, init_value=init_x_hc), + TensorSpec("query_start_loc", [2], torch.int32, init_value=torch.tensor([0, token_count], dtype=torch.int32)), TensorSpec("hc_attn_fn", [MIX_HC, HC_DIM], torch.float32, init_value=init_hc_attn_fn), TensorSpec("hc_attn_scale", [3], torch.float32, init_value=init_hc_attn_scale), TensorSpec("hc_attn_base", [MIX_HC], torch.float32, init_value=init_hc_attn_base), @@ -1037,7 +1075,7 @@ def init_wo_b(): ), TensorSpec( "compress_state_block_table", - [CSA_STATE_MAX_BLOCKS], + [1, CSA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table, ), @@ -1069,7 +1107,7 @@ def init_wo_b(): ), TensorSpec( "inner_compress_state_block_table", - [INNER_STATE_MAX_BLOCKS], + [1, INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table, ), @@ -1080,7 +1118,7 @@ def init_wo_b(): init_value=init_kv_cache, is_output=True, ), - TensorSpec("ori_block_table", [SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), + TensorSpec("ori_block_table", [1, SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), TensorSpec("ori_slot_mapping", [token_count], torch.int64, init_value=init_ori_slot_mapping), TensorSpec( "cmp_kv", @@ -1089,7 +1127,7 @@ def init_wo_b(): init_value=init_cmp_kv, is_output=True, ), - TensorSpec("cmp_block_table", [SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), + TensorSpec("cmp_block_table", [1, SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec( "idx_kv_cache", [IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], @@ -1104,8 +1142,9 @@ def init_wo_b(): init_value=init_idx_kv_scale, is_output=True, ), - TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), + TensorSpec("idx_block_table", [1, IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), + TensorSpec("local_request_ids", [token_count], torch.int32, init_value=init_local_request_ids), TensorSpec("cmp_slot_mapping", [token_count], torch.int64, init_value=init_cmp_slot_mapping), TensorSpec("idx_slot_mapping", [token_count], torch.int64, init_value=init_idx_slot_mapping), TensorSpec("state_slot_mapping", [token_count], torch.int64, init_value=init_state_slot_mapping), @@ -1138,6 +1177,7 @@ def _quant_w_per_output_channel_local(w): def prefill_attention_csa_cp_core( x_normed_local: pl.Tensor[[CP_Q_T_DYN, D], pl.BF16], x_normed_full: pl.Tensor[[CP_KV_T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], wq_a: pl.Tensor[[D, Q_LORA], pl.BF16], wq_b: pl.Tensor[[Q_LORA, H * HEAD_DIM], pl.INT8], wq_b_scale: pl.Tensor[[H * HEAD_DIM], pl.FP32], @@ -1157,7 +1197,7 @@ def prefill_attention_csa_cp_core( compress_state: pl.Tensor[ [MAIN_STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32 ], - compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], hadamard_idx: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], idx_wq_b_scale: pl.Tensor[[IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -1169,17 +1209,18 @@ def prefill_attention_csa_cp_core( inner_compress_state: pl.Tensor[ [INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32 ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8], idx_kv_scale: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], idx_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], state_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], @@ -1226,6 +1267,7 @@ def prefill_attention_csa_cp_core( compressor_ratio4( x_normed_full, + query_start_loc, compress_state, compress_state_block_table, cmp_wkv, cmp_wgate, cmp_ape, cmp_norm_w, cmp_freqs_cos_full, cmp_freqs_sin_full, @@ -1237,6 +1279,7 @@ def prefill_attention_csa_cp_core( indexer_cache_ready = pl.array.create(1, pl.TASK_ID) prefill_indexer_compressor( x_normed_full, + query_start_loc, inner_compress_state, inner_compress_state_block_table, inner_wkv, inner_wgate, inner_ape, inner_norm_w, cmp_freqs_cos_full, cmp_freqs_sin_full, @@ -1266,7 +1309,7 @@ def prefill_attention_csa_cp_core( hadamard_idx, idx_kv_cache, idx_kv_scale, idx_block_table, cmp_topk_indices, - position_ids_local, + position_ids_local, local_request_ids, indexer_cache_ready, ) @@ -1280,35 +1323,37 @@ def prefill_attention_csa_cp_core( swa_row = pl.full([1, WIN], dtype=pl.INT32, value=-1) mask_row = pl.full([1, VALID_BLOCK_MASK_COLS], dtype=pl.INT32, value=0) if t_idx < q_dim: - abs_pos = pl.read(position_ids_local, [t_idx]) - visible_cmp = pl.min((abs_pos + 1) // COMPRESS_RATIO, pl.cast(IDX_TOPK, pl.INT32)) - for mask_sb in pl.unroll(PREFILL_ATTN_BLOCKS): - cmp_lo = pl.max(mask_sb * PREFILL_ATTN_TILE - WIN, pl.cast(0, pl.INT32)) - cmp_hi_unclamped = (mask_sb + 1) * PREFILL_ATTN_TILE - WIN - cmp_hi_cap = pl.cast(SPARSE_CMP_BIAS_COLS, pl.INT32) - cmp_hi = pl.min(cmp_hi_unclamped, cmp_hi_cap) - if cmp_lo < cmp_hi: - if visible_cmp > cmp_lo: - pl.write(mask_row, [0, mask_sb], pl.cast(1, pl.INT32)) - window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) - key_start_abs = abs_pos + 1 - window_valid - for win_col in pl.range(WIN): - win_col_i32 = pl.cast(win_col, pl.INT32) - if win_col_i32 < window_valid: - key_abs = key_start_abs + win_col_i32 - blk_slot = key_abs // BLOCK_SIZE - blk = pl.read(ori_block_table, [pl.cast(blk_slot, pl.INDEX)]) - if blk >= 0: - row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) - pl.write(swa_row, [0, win_col], row) - pl.write(mask_row, [0, win_col // PREFILL_ATTN_TILE], pl.cast(1, pl.INT32)) + request_id = pl.read(local_request_ids, [t_idx]) + if request_id >= 0: + abs_pos = pl.read(position_ids_local, [t_idx]) + visible_cmp = pl.min((abs_pos + 1) // COMPRESS_RATIO, pl.cast(IDX_TOPK, pl.INT32)) + for mask_sb in pl.unroll(PREFILL_ATTN_BLOCKS): + cmp_lo = pl.max(mask_sb * PREFILL_ATTN_TILE - WIN, pl.cast(0, pl.INT32)) + cmp_hi_unclamped = (mask_sb + 1) * PREFILL_ATTN_TILE - WIN + cmp_hi_cap = pl.cast(SPARSE_CMP_BIAS_COLS, pl.INT32) + cmp_hi = pl.min(cmp_hi_unclamped, cmp_hi_cap) + if cmp_lo < cmp_hi: + if visible_cmp > cmp_lo: + pl.write(mask_row, [0, mask_sb], pl.cast(1, pl.INT32)) + window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) + key_start_abs = abs_pos + 1 - window_valid + for win_col in pl.range(WIN): + win_col_i32 = pl.cast(win_col, pl.INT32) + if win_col_i32 < window_valid: + key_abs = key_start_abs + win_col_i32 + blk_slot = key_abs // BLOCK_SIZE + blk = pl.read(ori_block_table, [request_id, pl.cast(blk_slot, pl.INDEX)]) + if blk >= 0: + row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) + pl.write(swa_row, [0, win_col], row) + pl.write(mask_row, [0, win_col // PREFILL_ATTN_TILE], pl.cast(1, pl.INT32)) swa_indices[t_idx : t_idx + 1, 0:WIN] = swa_row valid_block_mask[t_idx : t_idx + 1, 0:VALID_BLOCK_MASK_COLS] = mask_row attn_out_local = sparse_attn_physical( q, kv_cache, swa_indices, - cmp_kv, cmp_block_table, cmp_topk_indices, + cmp_kv, cmp_block_table, local_request_ids, cmp_topk_indices, valid_block_mask, attn_sink, freqs_cos_local, freqs_sin_local, wo_a, wo_b, wo_b_scale, @@ -1321,6 +1366,7 @@ def prefill_attention_csa_cp_core( @pl.jit.inline def prefill_attention_csa_cp( x_hc_full: pl.Tensor[[CP_KV_T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -1340,7 +1386,7 @@ def prefill_attention_csa_cp( cmp_ape: pl.Tensor[[COMPRESS_RATIO, MAIN_OUT_DIM], pl.FP32], cmp_norm_w: pl.Tensor[[HEAD_DIM], pl.BF16], compress_state: pl.Tensor[[MAIN_STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32], - compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], hadamard_idx: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], idx_wq_b_scale: pl.Tensor[[IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -1352,17 +1398,18 @@ def prefill_attention_csa_cp( inner_compress_state: pl.Tensor[ [INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32 ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8], idx_kv_scale: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], idx_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], state_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], @@ -1420,7 +1467,7 @@ def prefill_attention_csa_cp( attn_out_local = pl.create_tensor([q_dim, D], dtype=pl.BF16) attn_out_local = prefill_attention_csa_cp_core( - x_normed_local, x_normed_full, + x_normed_local, x_normed_full, query_start_loc, wq_a, wq_b, wq_b_scale, wkv, gamma_cq, gamma_ckv, freqs_cos_local, freqs_sin_local, @@ -1435,7 +1482,7 @@ def prefill_attention_csa_cp( kv_cache, ori_block_table, ori_slot_mapping_full, cmp_kv, cmp_block_table, idx_kv_cache, idx_kv_scale, idx_block_table, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, cmp_slot_mapping_full, idx_slot_mapping_full, state_slot_mapping_full, inner_state_slot_mapping_full, attn_sink, @@ -1458,6 +1505,7 @@ def prefill_attention_csa_cp( @pl.jit def prefill_attention_csa_cp_test( x_hc_full: pl.Tensor[[CP_KV_T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -1479,7 +1527,7 @@ def prefill_attention_csa_cp_test( compress_state: pl.InOut[ pl.Tensor[[MAIN_STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], hadamard_idx: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], idx_wq_b_scale: pl.Tensor[[IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -1491,17 +1539,18 @@ def prefill_attention_csa_cp_test( inner_compress_state: pl.InOut[ pl.Tensor[[INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32] ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], cmp_kv: pl.InOut[pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32]], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], idx_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], state_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], @@ -1522,18 +1571,25 @@ def prefill_attention_csa_cp_test( ): """Run one CP rank's query share with replicated CSA layer boundaries.""" x_hc_full.bind_dynamic(0, CP_KV_T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) freqs_cos.bind_dynamic(0, CP_KV_T_DYN) freqs_sin.bind_dynamic(0, CP_KV_T_DYN) cmp_freqs_cos.bind_dynamic(0, CP_KV_T_DYN) cmp_freqs_sin.bind_dynamic(0, CP_KV_T_DYN) compress_state.bind_dynamic(0, MAIN_STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) inner_compress_state.bind_dynamic(0, INNER_STATE_BLOCK_NUM_DYN) + inner_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) kv_cache.bind_dynamic(0, ORI_BLOCK_NUM_DYN) + ori_block_table.bind_dynamic(0, REQUESTS_DYN) cmp_kv.bind_dynamic(0, CMP_BLOCK_NUM_DYN) + cmp_block_table.bind_dynamic(0, REQUESTS_DYN) idx_kv_cache.bind_dynamic(0, IDX_BLOCK_NUM_DYN) idx_kv_scale.bind_dynamic(0, IDX_BLOCK_NUM_DYN) + idx_block_table.bind_dynamic(0, REQUESTS_DYN) ori_slot_mapping_full.bind_dynamic(0, CP_KV_T_DYN) position_ids_local.bind_dynamic(0, CP_Q_T_DYN) + local_request_ids.bind_dynamic(0, CP_Q_T_DYN) position_ids_full.bind_dynamic(0, CP_KV_T_DYN) cmp_slot_mapping_full.bind_dynamic(0, CP_KV_T_DYN) idx_slot_mapping_full.bind_dynamic(0, CP_KV_T_DYN) @@ -1547,6 +1603,7 @@ def prefill_attention_csa_cp_test( pl.write(o_proj_order_fence, [0], pl.cast(0, pl.INT32)) x_out_full, gather_signal = prefill_attention_csa_cp( x_hc_full, + query_start_loc, hc_attn_fn, hc_attn_scale, hc_attn_base, attn_norm_w, wq_a, wq_b, wq_b_scale, @@ -1563,7 +1620,7 @@ def prefill_attention_csa_cp_test( kv_cache, ori_block_table, ori_slot_mapping_full, cmp_kv, cmp_block_table, idx_kv_cache, idx_kv_scale, idx_block_table, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, cmp_slot_mapping_full, idx_slot_mapping_full, state_slot_mapping_full, inner_state_slot_mapping_full, attn_sink, @@ -1582,6 +1639,7 @@ def prefill_attention_csa_cp_test( @pl.jit.host def l3_prefill_attention_csa_cp( x_hc_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[TP_SIZE, QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[TP_SIZE, MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[TP_SIZE, 3], pl.FP32], hc_attn_base: pl.Tensor[[TP_SIZE, MIX_HC], pl.FP32], @@ -1603,7 +1661,7 @@ def l3_prefill_attention_csa_cp( compress_state: pl.InOut[ pl.Tensor[[TP_SIZE, MAIN_STATE_BLOCK_NUM_DYN, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[TP_SIZE, CSA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], hadamard_idx: pl.Tensor[[TP_SIZE, IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_wq_b: pl.Tensor[[TP_SIZE, Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], idx_wq_b_scale: pl.Tensor[[TP_SIZE, IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -1617,17 +1675,18 @@ def l3_prefill_attention_csa_cp( [TP_SIZE, INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32 ] ], - inner_compress_state_block_table: pl.Tensor[[TP_SIZE, INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[TP_SIZE, ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - ori_block_table: pl.Tensor[[TP_SIZE, SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT64], cmp_kv: pl.InOut[pl.Tensor[[TP_SIZE, CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - cmp_block_table: pl.Tensor[[TP_SIZE, SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.InOut[pl.Tensor[[TP_SIZE, IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[TP_SIZE, IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32]], - idx_block_table: pl.Tensor[[TP_SIZE, IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[TP_SIZE, CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[TP_SIZE, CP_Q_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT64], idx_slot_mapping_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT64], state_slot_mapping_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT64], @@ -1640,18 +1699,25 @@ def l3_prefill_attention_csa_cp( ): """Launch one CP group's CSA block, one child per rank.""" x_hc_full.bind_dynamic(1, CP_KV_T_DYN) + query_start_loc.bind_dynamic(1, QUERY_START_LOC_DYN) freqs_cos.bind_dynamic(1, CP_KV_T_DYN) freqs_sin.bind_dynamic(1, CP_KV_T_DYN) cmp_freqs_cos.bind_dynamic(1, CP_KV_T_DYN) cmp_freqs_sin.bind_dynamic(1, CP_KV_T_DYN) compress_state.bind_dynamic(1, MAIN_STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) inner_compress_state.bind_dynamic(1, INNER_STATE_BLOCK_NUM_DYN) + inner_compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) kv_cache.bind_dynamic(1, ORI_BLOCK_NUM_DYN) + ori_block_table.bind_dynamic(1, REQUESTS_DYN) cmp_kv.bind_dynamic(1, CMP_BLOCK_NUM_DYN) + cmp_block_table.bind_dynamic(1, REQUESTS_DYN) idx_kv_cache.bind_dynamic(1, IDX_BLOCK_NUM_DYN) idx_kv_scale.bind_dynamic(1, IDX_BLOCK_NUM_DYN) + idx_block_table.bind_dynamic(1, REQUESTS_DYN) ori_slot_mapping_full.bind_dynamic(1, CP_KV_T_DYN) position_ids_local.bind_dynamic(1, CP_Q_T_DYN) + local_request_ids.bind_dynamic(1, CP_Q_T_DYN) position_ids_full.bind_dynamic(1, CP_KV_T_DYN) cmp_slot_mapping_full.bind_dynamic(1, CP_KV_T_DYN) idx_slot_mapping_full.bind_dynamic(1, CP_KV_T_DYN) @@ -1678,6 +1744,7 @@ def l3_prefill_attention_csa_cp( o_proj_weight_consumed = pld.window(o_proj_weight_consumed_buf, [TP_SIZE, 1], dtype=pl.INT32) prefill_attention_csa_cp_test( x_hc_full[rank], + query_start_loc[rank], hc_attn_fn[rank], hc_attn_scale[rank], hc_attn_base[rank], attn_norm_w[rank], wq_a[rank], wq_b[rank], wq_b_scale[rank], @@ -1694,7 +1761,7 @@ def l3_prefill_attention_csa_cp( kv_cache[rank], ori_block_table[rank], ori_slot_mapping_full[rank], cmp_kv[rank], cmp_block_table[rank], idx_kv_cache[rank], idx_kv_scale[rank], idx_block_table[rank], - position_ids_local[rank], position_ids_full[rank], + position_ids_local[rank], position_ids_full[rank], local_request_ids[rank], cmp_slot_mapping_full[rank], idx_slot_mapping_full[rank], state_slot_mapping_full[rank], inner_state_slot_mapping_full[rank], attn_sink[rank], @@ -1756,6 +1823,11 @@ def build_cp_tensor_specs( "position_ids_full", [tp_size, token_count], spec.dtype, init_value=cp_stack(value, tp_size), )) + elif spec.name == "local_request_ids": + specs.append(TensorSpec( + "local_request_ids", [tp_size, local_t], spec.dtype, + init_value=value.reshape(tp_size, local_t).contiguous(), + )) elif spec.name == "wo_a": shards = [value[rank * O_PROJ_LOCAL_GROUPS : (rank + 1) * O_PROJ_LOCAL_GROUPS] for rank in range(tp_size)] specs.append(TensorSpec( @@ -1778,6 +1850,158 @@ def build_cp_tensor_specs( return specs +def build_ragged2_cp_tensor_specs(tp_size: int = TP_SIZE): + """Build the two-request rank-crossing CSA fixture from the B1 CP specs.""" + import torch + + from golden import TensorSpec + from prefill_cp_token_allgather import cp_stack + from utils import ( + block_table as make_block_table, + cache_row_from_table, + compressed_slot_mapping, + ori_slot_mapping as make_ori_slot_mapping, + state_slot_mapping as make_state_slot_mapping, + token_local_rope, + ) + + if tp_size != 2: + raise ValueError(f"ragged2 requires tp_size=2, got {tp_size}") + + token_count = 8 + request_starts = (126, 30) + request_positions = ( + torch.tensor([126, 127, 128], dtype=torch.int32), + torch.tensor([30, 31, 32, 33], dtype=torch.int32), + ) + position_ids = torch.cat((*request_positions, torch.zeros(1, dtype=torch.int32))) + query_start_loc = torch.tensor([0, 3, 7], dtype=torch.int32) + request_ids = torch.tensor([0, 0, 0, 1, 1, 1, 1, -1], dtype=torch.int32) + + ori_block_table = make_block_table(batch=2, table_blocks=SPARSE_ORI_MAX_BLOCKS, physical_blocks=CSA_ORI_BLOCK_NUM) + cmp_block_table = make_block_table(batch=2, table_blocks=SPARSE_CMP_MAX_BLOCKS, physical_blocks=CSA_CMP_BLOCK_NUM) + idx_block_table = make_block_table(batch=2, table_blocks=IDX_CACHE_MAX_BLOCKS, physical_blocks=IDX_CACHE_BLOCK_NUM) + compress_state_block_table = make_block_table( + batch=2, table_blocks=CSA_STATE_MAX_BLOCKS, + physical_blocks=CSA_STATE_BLOCK_NUM, + ) + inner_compress_state_block_table = make_block_table( + batch=2, table_blocks=INNER_STATE_MAX_BLOCKS, + physical_blocks=INNER_STATE_BLOCK_NUM, + ) + + ori_mappings = [] + cmp_mappings = [] + idx_mappings = [] + state_mappings = [] + inner_state_mappings = [] + state_size = CSA_STATE_BLOCK_SIZE + inner_state_size = INNER_STATE_BLOCK_SIZE + for request, positions in enumerate(request_positions): + positions_2d = positions.unsqueeze(0) + request_ori_table = ori_block_table[request : request + 1] + request_cmp_table = cmp_block_table[request : request + 1] + request_idx_table = idx_block_table[request : request + 1] + request_state_table = compress_state_block_table[request : request + 1] + inner_table = inner_compress_state_block_table[request : request + 1] + ori_mapping = make_ori_slot_mapping(positions_2d, request_ori_table) + cmp_mapping = compressed_slot_mapping(positions_2d, request_cmp_table, compress_ratio=COMPRESS_RATIO) + idx_mapping = compressed_slot_mapping(positions_2d, request_idx_table, compress_ratio=COMPRESS_RATIO) + state_mapping = make_state_slot_mapping(positions_2d, request_state_table, state_block_size=state_size) + inner_state_mapping = make_state_slot_mapping(positions_2d, inner_table, state_block_size=inner_state_size) + ori_mappings.append(ori_mapping.reshape(-1)) + cmp_mappings.append(cmp_mapping.reshape(-1)) + idx_mappings.append(idx_mapping.reshape(-1)) + state_mappings.append(state_mapping.reshape(-1)) + inner_state_mappings.append(inner_state_mapping.reshape(-1)) + pad_mapping = torch.full((1,), -1, dtype=torch.int64) + ori_slot_mapping = torch.cat((*ori_mappings, pad_mapping)) + cmp_slot_mapping = torch.cat((*cmp_mappings, pad_mapping)) + idx_slot_mapping = torch.cat((*idx_mappings, pad_mapping)) + state_slot_mapping = torch.cat((*state_mappings, pad_mapping)) + inner_state_slot_mapping = torch.cat((*inner_state_mappings, pad_mapping)) + + kv_cache = torch.zeros(CSA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.bfloat16) + kv_cache_flat = kv_cache.view(CSA_ORI_BLOCK_NUM * BLOCK_SIZE, HEAD_DIM) + for request, start_pos in enumerate(request_starts): + for position in range(max(0, start_pos - WIN), start_pos): + row = cache_row_from_table(ori_block_table[request], position) + kv_cache_flat[row] = ((torch.rand(HEAD_DIM) - 0.5) * 0.1).to(torch.bfloat16) + + cmp_kv = torch.zeros(CSA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.bfloat16) + idx_kv_cache = torch.zeros(IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM, dtype=torch.int8) + idx_kv_scale = torch.zeros(IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1, dtype=torch.float32) + compress_state_shape = (CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM) + inner_state_shape = (INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM) + compress_state = torch.zeros(compress_state_shape, dtype=torch.float32) + inner_compress_state = torch.zeros(inner_state_shape, dtype=torch.float32) + compress_state_flat = compress_state.view(-1, MAIN_COMPRESS_STATE_DIM) + inner_compress_state_flat = inner_compress_state.view(-1, INNER_COMPRESS_STATE_DIM) + for request, start_pos in enumerate(request_starts): + request_state_table = compress_state_block_table[request] + inner_table = inner_compress_state_block_table[request] + for position in range(max(0, start_pos - MAIN_STATE_LEN), start_pos): + row = cache_row_from_table(request_state_table, position, block_size=state_size) + compress_state_flat[row] = (torch.rand(MAIN_COMPRESS_STATE_DIM) - 0.5) * 0.05 + for position in range(max(0, start_pos - INNER_STATE_LEN), start_pos): + row = cache_row_from_table(inner_table, position, block_size=inner_state_size) + inner_compress_state_flat[row] = (torch.rand(INNER_COMPRESS_STATE_DIM) - 0.5) * 0.05 + + freqs_cos, freqs_sin = token_local_rope( + M, COMPRESS_RATIO, position_ids, + max_seq_len=MAX_SEQ_LEN, dtype=torch.bfloat16, + ) + cmp_positions = torch.where( + (position_ids + 1) % COMPRESS_RATIO == 0, + position_ids - (COMPRESS_RATIO - 1), + torch.zeros_like(position_ids), + ) + cmp_freqs_cos, cmp_freqs_sin = token_local_rope( + M, COMPRESS_RATIO, cmp_positions, + max_seq_len=MAX_SEQ_LEN, dtype=torch.bfloat16, + ) + + replacements = { + "query_start_loc": cp_stack(query_start_loc, tp_size), + "local_request_ids": request_ids.reshape(tp_size, token_count // tp_size).contiguous(), + "freqs_cos": cp_stack(freqs_cos, tp_size), + "freqs_sin": cp_stack(freqs_sin, tp_size), + "cmp_freqs_cos": cp_stack(cmp_freqs_cos, tp_size), + "cmp_freqs_sin": cp_stack(cmp_freqs_sin, tp_size), + "compress_state": cp_stack(compress_state, tp_size), + "compress_state_block_table": cp_stack(compress_state_block_table, tp_size), + "inner_compress_state": cp_stack(inner_compress_state, tp_size), + "inner_compress_state_block_table": cp_stack(inner_compress_state_block_table, tp_size), + "kv_cache": cp_stack(kv_cache, tp_size), + "ori_block_table": cp_stack(ori_block_table, tp_size), + "ori_slot_mapping_full": cp_stack(ori_slot_mapping, tp_size), + "cmp_kv": cp_stack(cmp_kv, tp_size), + "cmp_block_table": cp_stack(cmp_block_table, tp_size), + "idx_kv_cache": cp_stack(idx_kv_cache, tp_size), + "idx_kv_scale": cp_stack(idx_kv_scale, tp_size), + "idx_block_table": cp_stack(idx_block_table, tp_size), + "position_ids_local": position_ids.reshape(tp_size, token_count // tp_size).contiguous(), + "position_ids_full": cp_stack(position_ids, tp_size), + "cmp_slot_mapping_full": cp_stack(cmp_slot_mapping, tp_size), + "idx_slot_mapping_full": cp_stack(idx_slot_mapping, tp_size), + "state_slot_mapping_full": cp_stack(state_slot_mapping, tp_size), + "inner_state_slot_mapping_full": cp_stack(inner_state_slot_mapping, tp_size), + } + + specs = [] + for spec in build_cp_tensor_specs(start_pos=0, token_count=token_count, tp_size=tp_size): + value = replacements.get(spec.name) + if value is None: + specs.append(spec) + continue + replacement_spec = TensorSpec( + spec.name, list(value.shape), spec.dtype, init_value=value, + is_output=spec.is_output, resident=spec.resident, + ) + specs.append(replacement_spec) + return specs + + def golden_prefill_attention_csa_cp(tensors): """Run the single-die reference and replicate full outputs and caches across CP ranks.""" import torch @@ -1785,7 +2009,7 @@ def golden_prefill_attention_csa_cp(tensors): tp_size, token_count = tensors["x_hc_full"].shape[:2] shared = ( - "hc_attn_fn", "hc_attn_scale", "hc_attn_base", "attn_norm_w", + "query_start_loc", "hc_attn_fn", "hc_attn_scale", "hc_attn_base", "attn_norm_w", "wq_a", "wq_b", "wq_b_scale", "wkv", "gamma_cq", "gamma_ckv", "freqs_cos", "freqs_sin", "cmp_freqs_cos", "cmp_freqs_sin", "cmp_wkv", "cmp_wgate", "cmp_ape", "cmp_norm_w", @@ -1805,6 +2029,7 @@ def golden_prefill_attention_csa_cp(tensors): "state_slot_mapping", "inner_state_slot_mapping"): full[name] = tensors[f"{name}_full"][0] full["position_ids"] = tensors["position_ids_full"][0] + full["local_request_ids"] = tensors["local_request_ids"].reshape(token_count) full["x_out"] = torch.zeros(token_count, HC_MULT, D, dtype=torch.float32) golden_prefill_attention_csa(full) @@ -1836,8 +2061,12 @@ def golden_prefill_attention_csa_cp(tensors): parser.add_argument("--compile-only", action="store_true", default=False) parser.add_argument("--start-pos", type=int, default=START_POS) parser.add_argument( - "--token-count", "--num-tokens", dest="token_count", type=int, default=PREFILL_SEQ, - help="Physical query-token extent across the group; must divide by --tp.", + "--token-count", "--num-tokens", dest="token_count", type=int, default=None, + help=f"B1 physical query-token extent across the group; defaults to {PREFILL_SEQ}. ragged2 is fixed at 8.", + ) + parser.add_argument( + "--case", choices=["b1", "ragged2"], default="b1", + help="Fixture case; ragged2 is the fixed two-request TP2 boundary case.", ) parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) parser.add_argument("--enable-dep-gen", action="store_true", default=False) @@ -1845,7 +2074,7 @@ def golden_prefill_attention_csa_cp(tensors): args = parser.parse_args() # High-prefix sparse-attention tolerance. - x_out_diff_thd, x_out_max_diff = (8e-3, 2) if args.start_pos else (5e-3, 1) + x_out_diff_thd, x_out_max_diff = (8e-3, 2) if args.start_pos or args.case == "ragged2" else (5e-3, 1) cache_compare = { "kv_cache": ratio_allclose(atol=1e-4, rtol=1.0 / 128), "cmp_kv": ratio_allclose(atol=1e-4, rtol=1.0 / 128), @@ -1861,7 +2090,15 @@ def golden_prefill_attention_csa_cp(tensors): device_ids = [int(device) for device in args.device.split(",")] if len(device_ids) != TP_SIZE: parser.error(f"need exactly {TP_SIZE} devices, got {device_ids}") - if args.token_count % TP_SIZE != 0: + if args.case == "ragged2" and TP_SIZE != 2: + parser.error("--case ragged2 requires --tp 2") + if args.case == "ragged2" and args.start_pos != 0: + parser.error("--case ragged2 has fixed request starts and requires --start-pos 0") + if args.token_count is None: + args.token_count = 8 if args.case == "ragged2" else PREFILL_SEQ + if args.case == "ragged2" and args.token_count != 8: + parser.error("--case ragged2 has a fixed physical extent and requires --token-count 8") + if args.case == "b1" and args.token_count % TP_SIZE != 0: parser.error(f"--token-count must be a multiple of --tp={TP_SIZE}, got {args.token_count}") if TP_SIZE == 1: @@ -1893,9 +2130,14 @@ def golden_prefill_attention_csa_cp(tensors): else: from pypto.ir.distributed_compiled_program import DistributedConfig + specs = ( + build_ragged2_cp_tensor_specs(TP_SIZE) + if args.case == "ragged2" + else build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE) + ) result = run_jit( fn=l3_prefill_attention_csa_cp, - specs=build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE), + specs=specs, golden_fn=golden_prefill_attention_csa_cp, compile_cfg=dict( dump_passes=args.dump_passes, diff --git a/models/deepseek_v4_flash_dspark/prefill_fwd.py b/models/deepseek_v4_flash_dspark/prefill_fwd.py index 9ee36461e..4aed45a9f 100644 --- a/models/deepseek_v4_flash_dspark/prefill_fwd.py +++ b/models/deepseek_v4_flash_dspark/prefill_fwd.py @@ -41,6 +41,7 @@ from config import FLASH as MODEL_CONFIG from prefill_swa import ( build_cp_tensor_specs as build_swa_attention_tensor_specs, + build_ragged2_cp_tensor_specs as build_swa_ragged2_tensor_specs, prefill_attention_swa_cp, ) from prefill_hca import ( @@ -52,6 +53,7 @@ MAIN_OUT_DIM as HCA_MAIN_OUT_DIM, SPARSE_CMP_MAX_BLOCKS as HCA_CMP_MAX_BLOCKS, build_cp_tensor_specs as build_hca_attention_tensor_specs, + build_ragged2_cp_tensor_specs as build_hca_ragged2_tensor_specs, prefill_attention_hca_cp, ) from prefill_csa import ( @@ -82,12 +84,14 @@ SPARSE_ORI_MAX_BLOCKS, START_POS, build_cp_tensor_specs as build_csa_attention_tensor_specs, + build_ragged2_cp_tensor_specs as build_csa_ragged2_tensor_specs, prefill_attention_csa_cp, ) from prefill_cp_token_allgather import ( PREFILL_GROUP_CAP, TP_SIZE, ) +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN, lower_local_request_ids from prefill_o_proj import ( O_PROJ_LOCAL_COLS, O_PROJ_LOCAL_GROUPS, @@ -246,6 +250,7 @@ def mask_inactive_sample_rows( @pl.jit(auto_scope=False) def prefill_fwd( x_hc: pl.InOut[pl.Tensor[[FWD_GROUP_TOKENS_DYN, HC_MULT, D], pl.FP32]], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[FWD_NUM_LAYERS * MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[FWD_NUM_LAYERS * 3], pl.FP32], hc_attn_base: pl.Tensor[[FWD_NUM_LAYERS * MIX_HC], pl.FP32], @@ -284,9 +289,9 @@ def prefill_fwd( csa_inner_compress_state: pl.InOut[pl.Tensor[[FWD_INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, CSA_INNER_COMPRESS_STATE_DIM], pl.FP32]], idx_kv_cache: pl.InOut[pl.Tensor[[FWD_IDX_BLOCK_NUM_DYN, CSA_CMP_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[FWD_IDX_BLOCK_NUM_DYN, CSA_CMP_STORAGE_BLOCK_SIZE, 1, 1], pl.FP32]], - hca_compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], - csa_compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], - csa_inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + hca_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], + csa_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], + csa_inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], swa_freqs_cos: pl.Tensor[[FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], swa_freqs_sin: pl.Tensor[[FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], compressed_freqs_cos: pl.Tensor[[FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], @@ -295,10 +300,10 @@ def prefill_fwd( hca_cmp_freqs_sin: pl.Tensor[[FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], csa_cmp_freqs_cos: pl.Tensor[[FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], csa_cmp_freqs_sin: pl.Tensor[[FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], - hca_cmp_block_table: pl.Tensor[[HCA_CMP_MAX_BLOCKS], pl.INT32], - csa_cmp_block_table: pl.Tensor[[CSA_CMP_MAX_BLOCKS], pl.INT32], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], + hca_cmp_block_table: pl.Tensor[[REQUESTS_DYN, HCA_CMP_MAX_BLOCKS], pl.INT32], + csa_cmp_block_table: pl.Tensor[[REQUESTS_DYN, CSA_CMP_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[FWD_GROUP_TOKENS_DYN], pl.INT64], position_ids_local: pl.Tensor[[FWD_TOKENS_DYN], pl.INT32], position_ids_full: pl.Tensor[[FWD_GROUP_TOKENS_DYN], pl.INT32], @@ -367,6 +372,14 @@ def prefill_fwd( my_rank: pl.Scalar[pl.INT32], ): """Run the DeepSeek-V4 prefill backbone, LM head, and sampling.""" + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) + hca_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) + csa_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) + csa_inner_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) + ori_block_table.bind_dynamic(0, REQUESTS_DYN) + hca_cmp_block_table.bind_dynamic(0, REQUESTS_DYN) + csa_cmp_block_table.bind_dynamic(0, REQUESTS_DYN) + idx_block_table.bind_dynamic(0, REQUESTS_DYN) swa_freqs_cos.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) swa_freqs_sin.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) compressed_freqs_cos.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) @@ -377,7 +390,9 @@ def prefill_fwd( csa_cmp_freqs_sin.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) group_base = my_rank // TP_SIZE * TP_SIZE tp_rank = my_rank % TP_SIZE - + local_tokens = pl.tensor.dim(position_ids_local, 0) + local_request_ids = pl.create_tensor([local_tokens], dtype=pl.INT32) + lower_local_request_ids(query_start_loc, local_request_ids, tp_rank * local_tokens) ori_block_num = pl.tensor.dim(kv_cache, 0) // FWD_NUM_LAYERS hca_cmp_block_num = pl.tensor.dim(hca_cmp_kv, 0) // HCA_NUM_LAYERS csa_cmp_block_num = pl.tensor.dim(csa_cmp_kv, 0) // CSA_NUM_LAYERS @@ -452,7 +467,7 @@ def prefill_fwd( wkv_l0, gamma_cq_l0, gamma_ckv_l0, swa_freqs_cos, swa_freqs_sin, kv_cache_l0, ori_block_table, ori_slot_mapping_full, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, attn_sink_l0, wo_a_l0, wo_b_l0, wo_b_scale_l0, o_proj_wo_a_full, o_proj_wo_b_full, attn_stage, @@ -540,7 +555,7 @@ def prefill_fwd( wkv_l1, gamma_cq_l1, gamma_ckv_l1, swa_freqs_cos, swa_freqs_sin, kv_cache_l1, ori_block_table, ori_slot_mapping_full, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, attn_sink_l1, wo_a_l1, wo_b_l1, wo_b_scale_l1, o_proj_wo_a_full, o_proj_wo_b_full, attn_stage, @@ -677,6 +692,7 @@ def prefill_fwd( with pl.scope(): attn_stage, gather_signal = prefill_attention_csa_cp( x_hc, + query_start_loc, hc_attn_fn_csa, hc_attn_scale_csa, hc_attn_base_csa, attn_norm_w_csa, wq_a_csa, wq_b_csa, wq_b_scale_csa, wkv_csa, gamma_cq_csa, gamma_ckv_csa, @@ -691,7 +707,7 @@ def prefill_fwd( kv_cache_csa, ori_block_table, ori_slot_mapping_full, csa_cmp_kv_csa, csa_cmp_block_table, idx_kv_cache_csa, idx_kv_scale_csa, idx_block_table, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, csa_cmp_slot_mapping_full, csa_idx_slot_mapping_full, csa_state_slot_mapping_full, csa_inner_state_slot_mapping_full, attn_sink_csa, wo_a_csa, wo_b_csa, wo_b_scale_csa, @@ -795,6 +811,7 @@ def prefill_fwd( with pl.scope(): attn_stage, gather_signal = prefill_attention_hca_cp( x_hc, + query_start_loc, local_request_ids, hc_attn_fn_hca, hc_attn_scale_hca, hc_attn_base_hca, attn_norm_w_hca, wq_a_hca, wq_b_hca, wq_b_scale_hca, wkv_hca, gamma_cq_hca, gamma_ckv_hca, @@ -940,6 +957,7 @@ def prefill_fwd( with pl.scope(): attn_stage, gather_signal = prefill_attention_csa_cp( x_hc, + query_start_loc, hc_attn_fn_last, hc_attn_scale_last, hc_attn_base_last, attn_norm_w_last, wq_a_last, wq_b_last, wq_b_scale_last, wkv_last, gamma_cq_last, gamma_ckv_last, @@ -954,7 +972,7 @@ def prefill_fwd( kv_cache_last, ori_block_table, ori_slot_mapping_full, csa_cmp_kv_last, csa_cmp_block_table, idx_kv_cache_last, idx_kv_scale_last, idx_block_table, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, csa_cmp_slot_mapping_full, csa_idx_slot_mapping_full, csa_state_slot_mapping_full, csa_inner_state_slot_mapping_full, attn_sink_last, wo_a_last, wo_b_last, wo_b_scale_last, @@ -1014,6 +1032,7 @@ def prefill_fwd( @pl.jit.host def l3_prefill_fwd( x_hc: pl.InOut[pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, HC_MULT, D], pl.FP32]], + query_start_loc: pl.Tensor[[N_RANKS, QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * 3], pl.FP32], hc_attn_base: pl.Tensor[[N_RANKS, FWD_NUM_LAYERS * MIX_HC], pl.FP32], @@ -1052,9 +1071,9 @@ def l3_prefill_fwd( csa_inner_compress_state: pl.InOut[pl.Tensor[[N_RANKS, FWD_INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, CSA_INNER_COMPRESS_STATE_DIM], pl.FP32]], idx_kv_cache: pl.InOut[pl.Tensor[[N_RANKS, FWD_IDX_BLOCK_NUM_DYN, CSA_CMP_STORAGE_BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[N_RANKS, FWD_IDX_BLOCK_NUM_DYN, CSA_CMP_STORAGE_BLOCK_SIZE, 1, 1], pl.FP32]], - hca_compress_state_block_table: pl.Tensor[[N_RANKS, HCA_STATE_MAX_BLOCKS], pl.INT32], - csa_compress_state_block_table: pl.Tensor[[N_RANKS, CSA_STATE_MAX_BLOCKS], pl.INT32], - csa_inner_compress_state_block_table: pl.Tensor[[N_RANKS, INNER_STATE_MAX_BLOCKS], pl.INT32], + hca_compress_state_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], + csa_compress_state_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], + csa_inner_compress_state_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], swa_freqs_cos: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], swa_freqs_sin: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], compressed_freqs_cos: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], @@ -1063,10 +1082,10 @@ def l3_prefill_fwd( hca_cmp_freqs_sin: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], csa_cmp_freqs_cos: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], csa_cmp_freqs_sin: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, ROPE_HEAD_DIM], pl.BF16], - ori_block_table: pl.Tensor[[N_RANKS, SPARSE_ORI_MAX_BLOCKS], pl.INT32], - hca_cmp_block_table: pl.Tensor[[N_RANKS, HCA_CMP_MAX_BLOCKS], pl.INT32], - csa_cmp_block_table: pl.Tensor[[N_RANKS, CSA_CMP_MAX_BLOCKS], pl.INT32], - idx_block_table: pl.Tensor[[N_RANKS, IDX_CACHE_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], + hca_cmp_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, HCA_CMP_MAX_BLOCKS], pl.INT32], + csa_cmp_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, CSA_CMP_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN], pl.INT64], position_ids_local: pl.Tensor[[N_RANKS, FWD_TOKENS_DYN], pl.INT32], position_ids_full: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN], pl.INT32], @@ -1116,7 +1135,10 @@ def l3_prefill_fwd( ): """Run layer-major DSA-CP over a caller-padded physical token extent. - For logical length ``N``, each TP group supplies + Packed request boundaries are provided by a monotonic ``query_start_loc`` + that starts at zero and ends at total logical length ``N``. Its request + count must match the leading dimension of every request-indexed block table, + and metadata must be identical within a TP group. Each TP group supplies ``P = align_up(N, TP_SIZE)`` full rows and each rank supplies ``L = P // TP_SIZE`` local rows. Callers zero padded ``x_hc`` and ``input_ids``, use non-aliasing synthetic positions and ``-1`` cache/state @@ -1124,6 +1146,14 @@ def l3_prefill_fwd( The device schedule, including MoE collectives, runs over physical P/L. """ x_hc.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) + query_start_loc.bind_dynamic(1, QUERY_START_LOC_DYN) + hca_compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) + csa_compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) + csa_inner_compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) + ori_block_table.bind_dynamic(1, REQUESTS_DYN) + hca_cmp_block_table.bind_dynamic(1, REQUESTS_DYN) + csa_cmp_block_table.bind_dynamic(1, REQUESTS_DYN) + idx_block_table.bind_dynamic(1, REQUESTS_DYN) hidden_workspace.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) x_out.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) attn_stage.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) @@ -1196,6 +1226,7 @@ def l3_prefill_fwd( lm_head_logits_done = pld.window(lm_head_logits_done_buf, [LM_HEAD_TP_SIZE, 1], dtype=pl.INT32) prefill_fwd( x_hc[r], + query_start_loc[r], hc_attn_fn[r], hc_attn_scale[r], hc_attn_base[r], attn_norm_w[r], wq_a[r], wq_b[r], wq_b_scale[r], wkv[r], gamma_cq[r], gamma_ckv[r], @@ -1372,7 +1403,7 @@ def _global_token_index_map(local_tokens, torch): # Canonical host-tensor order for a single unified prefill layer. HOST_TENSOR_ORDER = ( - "x_hc", + "x_hc", "query_start_loc", "hc_attn_fn", "hc_attn_scale", "hc_attn_base", "attn_norm_w", "wq_a", "wq_b", "wq_b_scale", "wkv", "gamma_cq", "gamma_ckv", "swa_freqs_cos", "swa_freqs_sin", @@ -1425,21 +1456,36 @@ def _attention_kind_for_layer(layer_id): raise ValueError(f"unsupported DeepSeek V4 attention compress ratio {ratio} at layer {layer_id}") -def build_single_layer_tensor_specs(start_pos=START_POS, token_count=TP_SIZE * T, layer_id=2): +def build_single_layer_tensor_specs( + start_pos=START_POS, + token_count=TP_SIZE * T, + layer_id=2, + fixture_case="b1", +): """Build the single-layer tensor specs used by the stacked forward fixtures.""" import torch from golden import ScalarSpec, TensorSpec - def kind_specs(build_fn): + if fixture_case not in {"b1", "ragged2"}: + raise ValueError(f"unsupported full-forward fixture case {fixture_case!r}") + if fixture_case == "ragged2" and (TP_SIZE != 2 or token_count != 8): + raise ValueError(f"ragged2 requires TP=2 and physical token_count=8, got TP={TP_SIZE}, tokens={token_count}") + + def kind_specs(build_fn, build_ragged_fn): + source_specs = ( + build_ragged_fn(tp_size=TP_SIZE) + if fixture_case == "ragged2" + else build_fn(start_pos=start_pos, token_count=token_count, tp_size=TP_SIZE) + ) return { s.name: s - for s in build_fn(start_pos=start_pos, token_count=token_count, tp_size=TP_SIZE) + for s in source_specs if isinstance(s, TensorSpec) } - swa = kind_specs(build_swa_attention_tensor_specs) - hca = kind_specs(build_hca_attention_tensor_specs) - csa = kind_specs(build_csa_attention_tensor_specs) + swa = kind_specs(build_swa_attention_tensor_specs, build_swa_ragged2_tensor_specs) + hca = kind_specs(build_hca_attention_tensor_specs, build_hca_ragged2_tensor_specs) + csa = kind_specs(build_csa_attention_tensor_specs, build_csa_ragged2_tensor_specs) active_kind = _attention_kind_for_layer(layer_id) active = {"swa": swa, "hca": hca, "csa": csa}[active_kind] active_tokens = token_count // TP_SIZE @@ -1447,6 +1493,7 @@ def kind_specs(build_fn): # Unified names and source specs for the selected attention kind. attention_specs = [ ("x_hc", active["x_hc_full"]), + ("query_start_loc", csa["query_start_loc"]), ("hc_attn_fn", active["hc_attn_fn"]), ("hc_attn_scale", active["hc_attn_scale"]), ("hc_attn_base", active["hc_attn_base"]), ("attn_norm_w", active["attn_norm_w"]), @@ -1546,11 +1593,18 @@ def build_tensor_specs( hca_state_block_num=HCA_STATE_BLOCK_NUM, csa_state_block_num=CSA_STATE_BLOCK_NUM, inner_state_block_num=INNER_STATE_BLOCK_NUM, + fixture_case="b1", ): """Build CP-padded full-forward fixtures from a logical prompt length.""" import torch from golden import TensorSpec + if fixture_case not in {"b1", "ragged2"}: + raise ValueError(f"unsupported full-forward fixture case {fixture_case!r}") + if fixture_case == "ragged2" and TP_SIZE != 2: + raise ValueError(f"ragged2 requires TP=2, got TP={TP_SIZE}") + if fixture_case == "ragged2" and (start_pos != 0 or num_tokens != 7): + raise ValueError(f"ragged2 uses fixed start_pos=0 and logical num_tokens=7, got {start_pos} and {num_tokens}") if start_pos < 0: raise ValueError(f"start_pos must be non-negative, got {start_pos}") capacities = { @@ -1594,12 +1648,17 @@ def build_tensor_specs( base_specs = { spec.name: spec - for spec in build_single_layer_tensor_specs(start_pos=start_pos, token_count=physical_tokens, layer_id=0) + for spec in build_single_layer_tensor_specs( + start_pos=start_pos, + token_count=physical_tokens, + layer_id=0, + fixture_case=fixture_case, + ) if isinstance(spec, TensorSpec) } ordered_names = [ - "x_hc", + "x_hc", "query_start_loc", "hc_attn_fn", "hc_attn_scale", "hc_attn_base", "attn_norm_w", "wq_a", "wq_b", "wq_b_scale", "wkv", "gamma_cq", "gamma_ckv", "kv_cache", "attn_sink", "wo_a", "wo_b", "wo_b_scale", @@ -1648,7 +1707,21 @@ def build_tensor_specs( } specs = [] for name in ordered_names: - if name == "x_hc": + if name == "query_start_loc": + base = base_specs[name] + if fixture_case == "ragged2": + def init_query_start_loc(spec=base): + return _expand_rank_axis(_spec_value(spec, torch), torch) + + query_start_loc_shape = [N_RANKS, *base.shape[1:]] + else: + def init_query_start_loc(active_tokens=num_tokens, dtype=base.dtype): + boundaries = torch.tensor([0, active_tokens], dtype=dtype) + return boundaries.view(1, 2).expand(N_RANKS, -1).contiguous() + + query_start_loc_shape = [N_RANKS, 2] + specs.append(TensorSpec(name, query_start_loc_shape, base.dtype, init_value=init_query_start_loc)) + elif name == "x_hc": base = base_specs[name] x_hc_shape = list(base.shape) x_hc_shape[0] = N_RANKS @@ -1664,12 +1737,15 @@ def init_x_hc(tokens=physical_tokens, active_tokens=num_tokens, dtype=base.dtype specs.append(TensorSpec(name, x_hc_shape, base.dtype, init_value=init_x_hc, is_output=True)) elif name == "position_ids_local": - dtype = base_specs[name].dtype + if fixture_case == "ragged2": + specs.append(_make_shared_spec(name, base_specs)) + else: + dtype = base_specs[name].dtype - def init_position_ids_local(indices=global_token_indices, dtype=dtype): - return (start_pos + indices).to(dtype).contiguous() + def init_position_ids_local(indices=global_token_indices, dtype=dtype): + return (start_pos + indices).to(dtype).contiguous() - specs.append(TensorSpec(name, [N_RANKS, local_tokens], dtype, init_value=init_position_ids_local)) + specs.append(TensorSpec(name, [N_RANKS, local_tokens], dtype, init_value=init_position_ids_local)) elif name == "input_ids": dtype = base_specs[name].dtype @@ -1775,12 +1851,12 @@ def init_lm_head_weight(): shards = (torch.randn(TP_SIZE, VOCAB_PER_TP, D) / D**0.5).to(torch.bfloat16) return torch.stack([shards[rank % TP_SIZE] for rank in range(N_RANKS)], dim=0) - # Leader-owned rows: the group leader publishes the last prompt token as its - # single live logit row; peers keep every row at -1 and still join the TP - # collective. - def init_logit_row_indices(active_tokens=num_tokens): + # Group leaders publish one last-token row per packed request. + request_last_rows = (2, 6) if fixture_case == "ragged2" else (num_tokens - 1,) + + def init_logit_row_indices(last_rows=request_last_rows): indices = torch.full((N_RANKS, MAX_LOGIT_ROWS), -1, dtype=torch.int32) - indices[::TP_SIZE, 0] = active_tokens - 1 + indices[::TP_SIZE, : len(last_rows)] = torch.tensor(last_rows, dtype=torch.int32) return indices def init_hidden_workspace(): @@ -1801,6 +1877,24 @@ def init_hidden_workspace(): for spec in head_specs: spec.resident = "stacked" specs.append(spec) + + spec_by_name = {spec.name: spec for spec in specs} + request_count = spec_by_name["query_start_loc"].shape[1] - 1 + request_table_names = ( + "ori_block_table", "hca_cmp_block_table", "csa_cmp_block_table", "idx_block_table", + "hca_compress_state_block_table", "csa_compress_state_block_table", + "csa_inner_compress_state_block_table", + ) + mismatched = [ + f"{name}={spec_by_name[name].shape[1]}" + for name in request_table_names + if spec_by_name[name].shape[1] != request_count + ] + if mismatched: + raise ValueError( + f"request-indexed table rows must match query_start_loc request count {request_count}: " + f"{', '.join(mismatched)}" + ) return specs @@ -1941,8 +2035,12 @@ def main(): ) parser.add_argument("--start-pos", type=int, default=0) parser.add_argument( - "--num-tokens", type=int, default=TP_SIZE * T, - help="Prompt tokens per TP/CP group; must be at most 8192.", + "--num-tokens", type=int, default=None, + help=f"B1 prompt tokens per TP/CP group; defaults to {TP_SIZE * T}. ragged2 is fixed at 7.", + ) + parser.add_argument( + "--case", choices=["b1", "ragged2"], default="b1", + help="Fixture case; ragged2 is the fixed two-request TP2 boundary case.", ) parser.add_argument("--ori-block-num", type=int, default=CSA_ORI_BLOCK_NUM) parser.add_argument("--hca-cmp-block-num", type=int, default=HCA_CMP_BLOCK_NUM) @@ -1969,6 +2067,14 @@ def main(): parser.error(f"import-time N_RANKS must match --ep, got {N_RANKS} vs {args.ep}") if args.ep % args.tp != 0: parser.error(f"EP must be divisible by TP/CP, got --ep {args.ep} and --tp {args.tp}") + if args.case == "ragged2" and TP_SIZE != 2: + parser.error("--case ragged2 requires --tp 2") + if args.case == "ragged2" and args.start_pos != 0: + parser.error("--case ragged2 has fixed request starts and requires --start-pos 0") + if args.num_tokens is None: + args.num_tokens = 7 if args.case == "ragged2" else TP_SIZE * T + if args.case == "ragged2" and args.num_tokens != 7: + parser.error("--case ragged2 has a fixed logical extent and requires --num-tokens 7") if args.num_tokens < 1 or args.num_tokens > PREFILL_GROUP_CAP: parser.error(f"--num-tokens must be in [1, {PREFILL_GROUP_CAP}]") @@ -1982,6 +2088,7 @@ def main(): idx_block_num=args.idx_block_num, hca_state_block_num=args.hca_state_block_num, csa_state_block_num=args.csa_state_block_num, inner_state_block_num=args.inner_state_block_num, + fixture_case=args.case, ) result = run_jit( diff --git a/models/deepseek_v4_flash_dspark/prefill_hca.py b/models/deepseek_v4_flash_dspark/prefill_hca.py index 031318398..184c2f252 100644 --- a/models/deepseek_v4_flash_dspark/prefill_hca.py +++ b/models/deepseek_v4_flash_dspark/prefill_hca.py @@ -32,6 +32,7 @@ golden_prefill_compressor_ratio128, prefill_compressor_ratio128, ) +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN from qkv_proj_rope import golden_qkv_proj_rope, qkv_proj_rope from rmsnorm import golden_rms_norm, rms_norm from prefill_sparse_attn import ( @@ -108,6 +109,8 @@ @pl.jit.inline def prefill_attention_hca( x_hc: pl.Tensor[[T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -127,12 +130,12 @@ def prefill_attention_hca( cmp_ape: pl.Tensor[[COMPRESS_RATIO, MAIN_OUT_DIM], pl.FP32], cmp_norm_w: pl.Tensor[[HEAD_DIM], pl.BF16], compress_state: pl.Tensor[[STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32], - compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], ori_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], position_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], @@ -186,6 +189,7 @@ def prefill_attention_hca( prefill_compressor_ratio128( x_normed, + query_start_loc, compress_state, compress_state_block_table, cmp_wkv, @@ -204,18 +208,20 @@ def prefill_attention_hca( with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_hca_swa_indices") as swa_indices_tid: for idx_t in pl.range(t_dim): swa_row = pl.full([1, WIN], dtype=pl.INT32, value=-1) - abs_pos = pl.read(position_ids, [idx_t]) - window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) - key_start_abs = abs_pos + 1 - window_valid - for win_col in pl.range(WIN): - win_col_i32 = pl.cast(win_col, pl.INT32) - if win_col_i32 < window_valid: - key_abs = key_start_abs + win_col_i32 - blk_slot = key_abs // BLOCK_SIZE - blk = pl.read(ori_block_table, [pl.cast(blk_slot, pl.INDEX)]) - if blk >= 0: - row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) - pl.write(swa_row, [0, win_col], row) + request_id = pl.read(local_request_ids, [idx_t]) + if request_id >= 0: + abs_pos = pl.read(position_ids, [idx_t]) + window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) + key_start_abs = abs_pos + 1 - window_valid + for win_col in pl.range(WIN): + win_col_i32 = pl.cast(win_col, pl.INT32) + if win_col_i32 < window_valid: + key_abs = key_start_abs + win_col_i32 + blk_slot = key_abs // BLOCK_SIZE + blk = pl.read(ori_block_table, [request_id, pl.cast(blk_slot, pl.INDEX)]) + if blk >= 0: + row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) + pl.write(swa_row, [0, win_col], row) swa_indices[idx_t : idx_t + 1, 0:WIN] = swa_row # Streaming-attention input publication fence. @@ -253,17 +259,30 @@ def prefill_attention_hca( o_proj_weight_dep = pl.system.task_dummy(deps=[]) attn_out = pl.create_tensor([t_dim, D], dtype=pl.BF16) - attn_out = hca_streaming_attn_physical( - q, - kv_cache, swa_indices, - cmp_kv, cmp_block_table, - position_ids, attn_sink, - freqs_cos, freqs_sin, - wo_a, wo_b, wo_b_scale, - attn_out, - cache_ready_dep, - o_proj_weight_dep, - ) + with pl.spmd(t_dim, name_hint="prefill_hca_pad_output_init") as pad_output_tid: + pad_t = pl.tile.get_block_idx() + if pl.read(local_request_ids, [pad_t]) < 0: + attn_out[pad_t : pad_t + 1, :] = pl.full([1, D], dtype=pl.BF16, value=0.0) + request_dep = pl.system.task_dummy(deps=[cache_ready_dep, pad_output_tid]) + request_count = pl.tensor.dim(query_start_loc, 0) - 1 + for request in pl.range(request_count): + request_start = pl.cast(pl.read(query_start_loc, [request]), pl.INDEX) + request_end = pl.cast(pl.read(query_start_loc, [request + 1]), pl.INDEX) + request_rows = request_end - request_start + if request_rows > 0: + request_dep = hca_streaming_attn_physical( + q, + kv_cache, swa_indices, + cmp_kv, cmp_block_table[request], + position_ids, attn_sink, + freqs_cos, freqs_sin, + wo_a, wo_b, wo_b_scale, + attn_out, + request_dep, + o_proj_weight_dep, + request_start, + request_rows, + ) hc_post(attn_out, x_hc, post, comb, x_out) return x_out @@ -272,6 +291,8 @@ def prefill_attention_hca( @pl.jit def prefill_attention_hca_test( x_hc: pl.Tensor[[T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -293,12 +314,12 @@ def prefill_attention_hca_test( compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], ori_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], cmp_kv: pl.InOut[pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], position_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], @@ -309,10 +330,15 @@ def prefill_attention_hca_test( x_out: pl.Out[pl.Tensor[[T_DYN, HC_MULT, D], pl.FP32]], ): x_hc.bind_dynamic(0, T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) + local_request_ids.bind_dynamic(0, T_DYN) compress_state.bind_dynamic(0, STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) kv_cache.bind_dynamic(0, ORI_BLOCK_NUM_DYN) ori_slot_mapping.bind_dynamic(0, T_DYN) cmp_kv.bind_dynamic(0, CMP_BLOCK_NUM_DYN) + ori_block_table.bind_dynamic(0, REQUESTS_DYN) + cmp_block_table.bind_dynamic(0, REQUESTS_DYN) freqs_cos.bind_dynamic(0, T_DYN) freqs_sin.bind_dynamic(0, T_DYN) cmp_freqs_cos.bind_dynamic(0, T_DYN) @@ -324,6 +350,8 @@ def prefill_attention_hca_test( prefill_attention_hca( x_hc, + query_start_loc, + local_request_ids, hc_attn_fn, hc_attn_scale, hc_attn_base, @@ -428,37 +456,49 @@ def golden_prefill_attention_hca(tensors): ori_kv_flat[dst_row, :] = kv[t] cmp_kv = tensors["cmp_kv"] - golden_prefill_compressor_ratio128( - { - "x": x_normed.view(token_count, D), - "compress_state": tensors["compress_state"], - "compress_state_block_table": tensors["compress_state_block_table"], - "wkv": tensors["cmp_wkv"], - "wgate": tensors["cmp_wgate"], - "ape": tensors["cmp_ape"], - "norm_w": tensors["cmp_norm_w"], - "cmp_freqs_cos": tensors["cmp_freqs_cos"], - "cmp_freqs_sin": tensors["cmp_freqs_sin"], - "cmp_kv": cmp_kv, - "position_ids": tensors["position_ids"], - "cmp_slot_mapping": tensors["cmp_slot_mapping"], - "state_slot_mapping": tensors["state_slot_mapping"], - } - ) + query_start_loc = tensors["query_start_loc"] + for request in range(query_start_loc.numel() - 1): + request_start = int(query_start_loc[request].item()) + request_end = int(query_start_loc[request + 1].item()) + if request_end <= request_start: + continue + request_rows = slice(request_start, request_end) + golden_prefill_compressor_ratio128( + { + "x": x_normed[request_rows].view(request_end - request_start, D), + "compress_state": tensors["compress_state"], + "compress_state_block_table": tensors["compress_state_block_table"][request : request + 1], + "wkv": tensors["cmp_wkv"], + "wgate": tensors["cmp_wgate"], + "ape": tensors["cmp_ape"], + "norm_w": tensors["cmp_norm_w"], + "cmp_freqs_cos": tensors["cmp_freqs_cos"][request_rows], + "cmp_freqs_sin": tensors["cmp_freqs_sin"][request_rows], + "cmp_kv": cmp_kv, + "position_ids": tensors["position_ids"][request_rows], + "cmp_slot_mapping": tensors["cmp_slot_mapping"][request_rows], + "state_slot_mapping": tensors["state_slot_mapping"][request_rows], + } + ) def build_sparse_metadata(): swa_idx = torch.full((token_count, WIN), -1, dtype=torch.int32) pos = tensors["position_ids"] - max_visible_cmp = min(int((pos[-1].item() + 1) // COMPRESS_RATIO), SPARSE_CMP_MAX_BLOCKS * BLOCK_SIZE) + request_ids = tensors["local_request_ids"] + active = request_ids >= 0 + max_position = int(pos[active].max().item()) if active.any() else -1 + max_visible_cmp = min((max_position + 1) // COMPRESS_RATIO, SPARSE_CMP_MAX_BLOCKS * BLOCK_SIZE) cmp_idx = torch.full((token_count, max(1, max_visible_cmp)), -1, dtype=torch.int32) - ori_table = tensors["ori_block_table"] cmp_cap = SPARSE_CMP_MAX_BLOCKS * BLOCK_SIZE for t in range(token_count): + request_id = int(request_ids[t].item()) + if request_id < 0: + continue abs_pos = int(pos[t].item()) window_valid = min(WIN, abs_pos + 1) key_start_abs = abs_pos + 1 - window_valid for k, key_abs in enumerate(range(key_start_abs, abs_pos + 1)): - row = cache_row_from_table(ori_table, key_abs) + row = cache_row_from_table(tensors["ori_block_table"][request_id], key_abs) if row >= 0: swa_idx[t, k] = row visible_cmp = min((abs_pos + 1) // COMPRESS_RATIO, max_visible_cmp, cmp_cap) @@ -475,6 +515,7 @@ def build_sparse_metadata(): "swa_indices": swa_indices, "cmp_kv": cmp_kv, "cmp_block_table": tensors["cmp_block_table"], + "local_request_ids": tensors["local_request_ids"], "cmp_indices": cmp_indices, "attn_sink": tensors["attn_sink"], "freqs_cos": rope_cos_t, @@ -557,6 +598,12 @@ def cmp_write_records(): def init_x_hc(): return torch.empty(token_count, HC_MULT, D).uniform_(-1, 1) + def init_query_start_loc(): + return torch.tensor([0, token_count], dtype=torch.int32) + + def init_local_request_ids(): + return torch.zeros(token_count, dtype=torch.int32) + # Real layer-9 (HCA, ratio-128) hc_attn scale/base, fn synthetic at real magnitude. A synthetic # scale=0.5/base=0 cancels attn_out and the hc residual to near-zero in x_out, where W8A8 noise # blows up the relative tail. @@ -649,7 +696,7 @@ def init_cmp_norm_w(): state_table = _state_block_table(HCA_STATE_MAX_BLOCKS, HCA_STATE_PHYSICAL_BLOCKS) def init_compress_state_block_table(): - return state_table.clone() + return state_table.clone().unsqueeze(0) def state_row(abs_pos): if abs_pos < 0 or abs_pos >= MAX_SEQ_LEN: @@ -675,7 +722,7 @@ def init_compress_state(): def init_kv_cache(): cache = torch.zeros(HCA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM) cache_flat = cache.view(HCA_ORI_BLOCK_NUM * BLOCK_SIZE, HEAD_DIM) - table = init_ori_block_table() + table = init_ori_block_table()[0] if context_len > 0: prefix_start = max(0, context_len - WIN) prefix = ((torch.rand(context_len - prefix_start, HEAD_DIM) - 0.5) * 0.1).to(torch.bfloat16) @@ -688,7 +735,7 @@ def init_kv_cache(): def init_ori_slot_mapping(): mapping = torch.full((token_count,), -1, dtype=torch.int64) local_pos, _ = token_meta() - table = init_ori_block_table() + table = init_ori_block_table()[0] for t in range(token_count): logical_pos = context_len + int(local_pos[t].item()) mapping[t] = cache_row_from_table(table, logical_pos) @@ -698,12 +745,12 @@ def init_ori_block_table(): table = torch.full((SPARSE_ORI_MAX_BLOCKS,), -1, dtype=torch.int32) for block in range(SPARSE_ORI_MAX_BLOCKS): table[block] = block % HCA_ORI_BLOCK_NUM - return table + return table.unsqueeze(0) def init_cmp_kv(): cache = torch.zeros(HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM) cache_flat = cache.view(HCA_CMP_BLOCK_NUM * BLOCK_SIZE, HEAD_DIM) - table = init_cmp_block_table() + table = init_cmp_block_table()[0] completed = context_len // COMPRESS_RATIO if completed > 0: prefix_cmp = ((torch.rand(completed, HEAD_DIM) - 0.5) * 0.1).to(torch.bfloat16) @@ -717,14 +764,14 @@ def init_cmp_block_table(): table = torch.full((SPARSE_CMP_MAX_BLOCKS,), -1, dtype=torch.int32) for block in range(min(SPARSE_CMP_MAX_BLOCKS, HCA_CMP_BLOCK_NUM)): table[block] = block - return table + return table.unsqueeze(0) def init_position_ids(): return token_meta()[1] def init_cmp_slot_mapping(): out = torch.full((token_count,), -1, dtype=torch.int64) - table = init_cmp_block_table() + table = init_cmp_block_table()[0] records = cmp_write_records() for token_id, cmp_slot in records: out[token_id] = cache_row_from_table(table, cmp_slot) @@ -753,6 +800,8 @@ def init_wo_b(): return [ TensorSpec("x_hc", [token_count, HC_MULT, D], torch.float32, init_value=init_x_hc), + TensorSpec("query_start_loc", [2], torch.int32, init_value=init_query_start_loc), + TensorSpec("local_request_ids", [token_count], torch.int32, init_value=init_local_request_ids), TensorSpec("hc_attn_fn", [MIX_HC, HC_DIM], torch.float32, init_value=init_hc_attn_fn), TensorSpec("hc_attn_scale", [3], torch.float32, init_value=init_hc_attn_scale), TensorSpec("hc_attn_base", [MIX_HC], torch.float32, init_value=init_hc_attn_base), @@ -780,7 +829,7 @@ def init_wo_b(): ), TensorSpec( "compress_state_block_table", - [HCA_STATE_MAX_BLOCKS], + [1, HCA_STATE_MAX_BLOCKS], torch.int32, init_value=init_compress_state_block_table, ), @@ -792,7 +841,7 @@ def init_wo_b(): is_output=True, ), TensorSpec("ori_slot_mapping", [token_count], torch.int64, init_value=init_ori_slot_mapping), - TensorSpec("ori_block_table", [SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), + TensorSpec("ori_block_table", [1, SPARSE_ORI_MAX_BLOCKS], torch.int32, init_value=init_ori_block_table), TensorSpec( "cmp_kv", [HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], @@ -800,7 +849,7 @@ def init_wo_b(): init_value=init_cmp_kv, is_output=True, ), - TensorSpec("cmp_block_table", [SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), + TensorSpec("cmp_block_table", [1, SPARSE_CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), TensorSpec("cmp_slot_mapping", [token_count], torch.int64, init_value=init_cmp_slot_mapping), TensorSpec("state_slot_mapping", [token_count], torch.int64, init_value=init_state_slot_mapping), @@ -816,6 +865,8 @@ def init_wo_b(): def prefill_attention_hca_cp_core( x_normed_local: pl.Tensor[[CP_Q_T_DYN, D], pl.BF16], x_normed_full: pl.Tensor[[CP_KV_T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], wq_a: pl.Tensor[[D, Q_LORA], pl.BF16], wq_b: pl.Tensor[[Q_LORA, H * HEAD_DIM], pl.INT8], wq_b_scale: pl.Tensor[[H * HEAD_DIM], pl.FP32], @@ -833,12 +884,12 @@ def prefill_attention_hca_cp_core( cmp_ape: pl.Tensor[[COMPRESS_RATIO, MAIN_OUT_DIM], pl.FP32], cmp_norm_w: pl.Tensor[[HEAD_DIM], pl.BF16], compress_state: pl.Tensor[[STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32], - compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], @@ -885,6 +936,7 @@ def prefill_attention_hca_cp_core( prefill_compressor_ratio128( x_normed_full, + query_start_loc, compress_state, compress_state_block_table, cmp_wkv, cmp_wgate, cmp_ape, cmp_norm_w, cmp_freqs_cos_full, cmp_freqs_sin_full, @@ -896,18 +948,20 @@ def prefill_attention_hca_cp_core( with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_hca_cp_swa_indices") as swa_indices_tid: for idx_t in pl.range(q_dim): swa_row = pl.full([1, WIN], dtype=pl.INT32, value=-1) - abs_pos = pl.read(position_ids_local, [idx_t]) - window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) - key_start_abs = abs_pos + 1 - window_valid - for win_col in pl.range(WIN): - win_col_i32 = pl.cast(win_col, pl.INT32) - if win_col_i32 < window_valid: - key_abs = key_start_abs + win_col_i32 - blk_slot = key_abs // BLOCK_SIZE - blk = pl.read(ori_block_table, [pl.cast(blk_slot, pl.INDEX)]) - if blk >= 0: - row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) - pl.write(swa_row, [0, win_col], row) + request_id = pl.read(local_request_ids, [idx_t]) + if request_id >= 0: + abs_pos = pl.read(position_ids_local, [idx_t]) + window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) + key_start_abs = abs_pos + 1 - window_valid + for win_col in pl.range(WIN): + win_col_i32 = pl.cast(win_col, pl.INT32) + if win_col_i32 < window_valid: + key_abs = key_start_abs + win_col_i32 + blk_slot = key_abs // BLOCK_SIZE + blk = pl.read(ori_block_table, [request_id, pl.cast(blk_slot, pl.INDEX)]) + if blk >= 0: + row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) + pl.write(swa_row, [0, win_col], row) swa_indices[idx_t : idx_t + 1, 0:WIN] = swa_row # Streaming-attention input publication fence. @@ -944,23 +998,44 @@ def prefill_attention_hca_cp_core( ready_bit = ready_bit + cmp_ready_value pl.write(cache_ready_fence, [0], ready_bit) - attn_out_local = hca_streaming_attn_physical( - q, - kv_cache, swa_indices, - cmp_kv, cmp_block_table, - position_ids_local, attn_sink, - freqs_cos_local, freqs_sin_local, - wo_a, wo_b, wo_b_scale, - attn_out_local, - cache_ready_dep, - o_proj_weight_dep, - ) + # Per-request HCA streaming over rank-local packed query intervals. + with pl.spmd(q_dim, name_hint="prefill_hca_cp_pad_output_init") as pad_output_tid: + pad_t = pl.tile.get_block_idx() + if pl.read(local_request_ids, [pad_t]) < 0: + attn_out_local[pad_t : pad_t + 1, :] = pl.full([1, D], dtype=pl.BF16, value=0.0) + request_dep = pl.system.task_dummy(deps=[cache_ready_dep, pad_output_tid]) + request_count = pl.tensor.dim(query_start_loc, 0) - 1 + for request in pl.range(request_count): + local_start = pl.cast(0, pl.INDEX) + request_rows = pl.cast(0, pl.INDEX) + for local_t in pl.range(q_dim): + request_id = pl.read(local_request_ids, [local_t]) + if request_id == request: + if request_rows == 0: + local_start = local_t + request_rows = request_rows + 1 + if request_rows > 0: + request_dep = hca_streaming_attn_physical( + q, + kv_cache, swa_indices, + cmp_kv, cmp_block_table[request], + position_ids_local, attn_sink, + freqs_cos_local, freqs_sin_local, + wo_a, wo_b, wo_b_scale, + attn_out_local, + request_dep, + o_proj_weight_dep, + local_start, + request_rows, + ) return attn_out_local @pl.jit.inline def prefill_attention_hca_cp( x_hc_full: pl.Tensor[[CP_KV_T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -980,12 +1055,12 @@ def prefill_attention_hca_cp( cmp_ape: pl.Tensor[[COMPRESS_RATIO, MAIN_OUT_DIM], pl.FP32], cmp_norm_w: pl.Tensor[[HEAD_DIM], pl.BF16], compress_state: pl.Tensor[[STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32], - compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], @@ -1049,6 +1124,7 @@ def prefill_attention_hca_cp( attn_out_local = pl.create_tensor([q_dim, D], dtype=pl.BF16) attn_out_local = prefill_attention_hca_cp_core( x_normed_local, x_normed_full, + query_start_loc, local_request_ids, wq_a, wq_b, wq_b_scale, wkv, gamma_cq, gamma_ckv, freqs_cos_local, freqs_sin_local, @@ -1080,6 +1156,8 @@ def prefill_attention_hca_cp( @pl.jit def prefill_attention_hca_cp_test( x_hc_full: pl.Tensor[[CP_KV_T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -1101,12 +1179,12 @@ def prefill_attention_hca_cp_test( compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], cmp_kv: pl.InOut[pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - cmp_block_table: pl.Tensor[[SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], @@ -1127,9 +1205,14 @@ def prefill_attention_hca_cp_test( ): """Run one DSA-CP rank's share of an HCA block.""" x_hc_full.bind_dynamic(0, CP_KV_T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) + local_request_ids.bind_dynamic(0, CP_Q_T_DYN) compress_state.bind_dynamic(0, STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) kv_cache.bind_dynamic(0, ORI_BLOCK_NUM_DYN) cmp_kv.bind_dynamic(0, CMP_BLOCK_NUM_DYN) + ori_block_table.bind_dynamic(0, REQUESTS_DYN) + cmp_block_table.bind_dynamic(0, REQUESTS_DYN) ori_slot_mapping_full.bind_dynamic(0, CP_KV_T_DYN) freqs_cos.bind_dynamic(0, CP_KV_T_DYN) freqs_sin.bind_dynamic(0, CP_KV_T_DYN) @@ -1148,6 +1231,7 @@ def prefill_attention_hca_cp_test( pl.write(o_proj_order_fence, [0], pl.cast(0, pl.INT32)) x_out_full, gather_signal = prefill_attention_hca_cp( x_hc_full, + query_start_loc, local_request_ids, hc_attn_fn, hc_attn_scale, hc_attn_base, attn_norm_w, wq_a, wq_b, wq_b_scale, wkv, gamma_cq, gamma_ckv, freqs_cos, freqs_sin, @@ -1173,6 +1257,8 @@ def prefill_attention_hca_cp_test( @pl.jit.host def l3_prefill_attention_hca_cp( x_hc_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[TP_SIZE, QUERY_START_LOC_DYN], pl.INT32], + local_request_ids: pl.Tensor[[TP_SIZE, CP_Q_T_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[TP_SIZE, MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[TP_SIZE, 3], pl.FP32], hc_attn_base: pl.Tensor[[TP_SIZE, MIX_HC], pl.FP32], @@ -1194,12 +1280,12 @@ def l3_prefill_attention_hca_cp( compress_state: pl.InOut[ pl.Tensor[[TP_SIZE, STATE_BLOCK_NUM_DYN, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM], pl.FP32] ], - compress_state_block_table: pl.Tensor[[TP_SIZE, HCA_STATE_MAX_BLOCKS], pl.INT32], + compress_state_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[TP_SIZE, ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], ori_slot_mapping_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT64], - ori_block_table: pl.Tensor[[TP_SIZE, SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], cmp_kv: pl.InOut[pl.Tensor[[TP_SIZE, CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - cmp_block_table: pl.Tensor[[TP_SIZE, SPARSE_CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, SPARSE_CMP_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[TP_SIZE, CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT32], cmp_slot_mapping_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT64], @@ -1212,9 +1298,14 @@ def l3_prefill_attention_hca_cp( ): """Launch one DSA-CP HCA block per rank.""" x_hc_full.bind_dynamic(1, CP_KV_T_DYN) + query_start_loc.bind_dynamic(1, QUERY_START_LOC_DYN) + local_request_ids.bind_dynamic(1, CP_Q_T_DYN) compress_state.bind_dynamic(1, STATE_BLOCK_NUM_DYN) + compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) kv_cache.bind_dynamic(1, ORI_BLOCK_NUM_DYN) cmp_kv.bind_dynamic(1, CMP_BLOCK_NUM_DYN) + ori_block_table.bind_dynamic(1, REQUESTS_DYN) + cmp_block_table.bind_dynamic(1, REQUESTS_DYN) ori_slot_mapping_full.bind_dynamic(1, CP_KV_T_DYN) freqs_cos.bind_dynamic(1, CP_KV_T_DYN) freqs_sin.bind_dynamic(1, CP_KV_T_DYN) @@ -1246,6 +1337,7 @@ def l3_prefill_attention_hca_cp( o_proj_weight_consumed = pld.window(o_proj_weight_consumed_buf, [TP_SIZE, 1], dtype=pl.INT32) prefill_attention_hca_cp_test( x_hc_full[rank], + query_start_loc[rank], local_request_ids[rank], hc_attn_fn[rank], hc_attn_scale[rank], hc_attn_base[rank], attn_norm_w[rank], wq_a[rank], wq_b[rank], wq_b_scale[rank], wkv[rank], gamma_cq[rank], gamma_ckv[rank], @@ -1295,7 +1387,17 @@ def build_cp_tensor_specs( specs = [] for spec in build_tensor_specs(start_pos, token_count): value = materialize_spec(spec) - if spec.name == "x_hc": + if spec.name == "query_start_loc": + specs.append(TensorSpec( + "query_start_loc", [tp_size, 2], spec.dtype, + init_value=cp_stack(value, tp_size), + )) + elif spec.name == "local_request_ids": + specs.append(TensorSpec( + "local_request_ids", [tp_size, local_t], spec.dtype, + init_value=value.reshape(tp_size, local_t).contiguous(), + )) + elif spec.name == "x_hc": specs.append(TensorSpec( "x_hc_full", [tp_size, token_count, HC_MULT, D], spec.dtype, init_value=cp_stack(value, tp_size), @@ -1336,6 +1438,126 @@ def build_cp_tensor_specs( return specs +def build_ragged2_cp_tensor_specs(tp_size: int = TP_SIZE): + """Build the two-request rank-crossing HCA fixture from the B1 CP specs.""" + import torch + + from golden import TensorSpec + from prefill_cp_token_allgather import cp_stack + from utils import ( + block_table as make_block_table, + cache_row_from_table, + compressed_slot_mapping, + ori_slot_mapping as make_ori_slot_mapping, + state_slot_mapping as make_state_slot_mapping, + token_local_rope, + ) + + if tp_size != 2: + raise ValueError(f"ragged2 requires tp_size=2, got {tp_size}") + + token_count = 8 + request_starts = (126, 30) + request_positions = ( + torch.tensor([126, 127, 128], dtype=torch.int32), + torch.tensor([30, 31, 32, 33], dtype=torch.int32), + ) + position_ids = torch.cat((*request_positions, torch.zeros(1, dtype=torch.int32))) + query_start_loc = torch.tensor([0, 3, 7], dtype=torch.int32) + request_ids = torch.tensor([0, 0, 0, 1, 1, 1, 1, -1], dtype=torch.int32) + + ori_block_table = make_block_table(batch=2, table_blocks=SPARSE_ORI_MAX_BLOCKS, physical_blocks=HCA_ORI_BLOCK_NUM) + cmp_block_table = make_block_table(batch=2, table_blocks=SPARSE_CMP_MAX_BLOCKS, physical_blocks=HCA_CMP_BLOCK_NUM) + compress_state_block_table = make_block_table( + batch=2, table_blocks=HCA_STATE_MAX_BLOCKS, + physical_blocks=HCA_STATE_BLOCK_NUM, + ) + + ori_mappings = [] + cmp_mappings = [] + state_mappings = [] + state_size = HCA_STATE_BLOCK_SIZE + for request, positions in enumerate(request_positions): + positions_2d = positions.unsqueeze(0) + request_ori_table = ori_block_table[request : request + 1] + request_cmp_table = cmp_block_table[request : request + 1] + request_state_table = compress_state_block_table[request : request + 1] + ori_mapping = make_ori_slot_mapping(positions_2d, request_ori_table) + cmp_mapping = compressed_slot_mapping(positions_2d, request_cmp_table, compress_ratio=COMPRESS_RATIO) + state_mapping = make_state_slot_mapping(positions_2d, request_state_table, state_block_size=state_size) + ori_mappings.append(ori_mapping.reshape(-1)) + cmp_mappings.append(cmp_mapping.reshape(-1)) + state_mappings.append(state_mapping.reshape(-1)) + pad_mapping = torch.full((1,), -1, dtype=torch.int64) + ori_slot_mapping = torch.cat((*ori_mappings, pad_mapping)) + cmp_slot_mapping = torch.cat((*cmp_mappings, pad_mapping)) + state_slot_mapping = torch.cat((*state_mappings, pad_mapping)) + + kv_cache = torch.zeros(HCA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.bfloat16) + kv_cache_flat = kv_cache.view(HCA_ORI_BLOCK_NUM * BLOCK_SIZE, HEAD_DIM) + for request, start_pos in enumerate(request_starts): + for position in range(max(0, start_pos - WIN), start_pos): + row = cache_row_from_table(ori_block_table[request], position) + kv_cache_flat[row] = ((torch.rand(HEAD_DIM) - 0.5) * 0.1).to(torch.bfloat16) + + cmp_kv = torch.zeros(HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.bfloat16) + compress_state_shape = (HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, MAIN_COMPRESS_STATE_DIM) + compress_state = torch.zeros(compress_state_shape, dtype=torch.float32) + compress_state_flat = compress_state.view(-1, MAIN_COMPRESS_STATE_DIM) + for request, start_pos in enumerate(request_starts): + request_state_table = compress_state_block_table[request] + for position in range(max(0, start_pos - COMPRESS_RATIO), start_pos): + row = cache_row_from_table(request_state_table, position, block_size=state_size) + compress_state_flat[row] = (torch.rand(MAIN_COMPRESS_STATE_DIM) - 0.5) * 0.05 + + freqs_cos, freqs_sin = token_local_rope( + M, COMPRESS_RATIO, position_ids, + max_seq_len=MAX_SEQ_LEN, dtype=torch.bfloat16, + ) + cmp_positions = torch.where( + (position_ids + 1) % COMPRESS_RATIO == 0, + position_ids - (COMPRESS_RATIO - 1), + torch.zeros_like(position_ids), + ) + cmp_freqs_cos, cmp_freqs_sin = token_local_rope( + M, COMPRESS_RATIO, cmp_positions, + max_seq_len=MAX_SEQ_LEN, dtype=torch.bfloat16, + ) + + replacements = { + "query_start_loc": cp_stack(query_start_loc, tp_size), + "local_request_ids": request_ids.reshape(tp_size, token_count // tp_size).contiguous(), + "freqs_cos": cp_stack(freqs_cos, tp_size), + "freqs_sin": cp_stack(freqs_sin, tp_size), + "cmp_freqs_cos": cp_stack(cmp_freqs_cos, tp_size), + "cmp_freqs_sin": cp_stack(cmp_freqs_sin, tp_size), + "compress_state": cp_stack(compress_state, tp_size), + "compress_state_block_table": cp_stack(compress_state_block_table, tp_size), + "kv_cache": cp_stack(kv_cache, tp_size), + "ori_slot_mapping_full": cp_stack(ori_slot_mapping, tp_size), + "ori_block_table": cp_stack(ori_block_table, tp_size), + "cmp_kv": cp_stack(cmp_kv, tp_size), + "cmp_block_table": cp_stack(cmp_block_table, tp_size), + "position_ids_local": position_ids.reshape(tp_size, token_count // tp_size).contiguous(), + "position_ids_full": cp_stack(position_ids, tp_size), + "cmp_slot_mapping_full": cp_stack(cmp_slot_mapping, tp_size), + "state_slot_mapping_full": cp_stack(state_slot_mapping, tp_size), + } + + specs = [] + for spec in build_cp_tensor_specs(start_pos=0, token_count=token_count, tp_size=tp_size): + value = replacements.get(spec.name) + if value is None: + specs.append(spec) + continue + replacement_spec = TensorSpec( + spec.name, list(value.shape), spec.dtype, init_value=value, + is_output=spec.is_output, resident=spec.resident, + ) + specs.append(replacement_spec) + return specs + + def golden_prefill_attention_hca_cp(tensors): """Run the full-stream reference and replicate layer outputs and caches per rank.""" import torch @@ -1361,6 +1583,10 @@ def golden_prefill_attention_hca_cp(tensors): full["cmp_slot_mapping"] = tensors["cmp_slot_mapping_full"][0] full["state_slot_mapping"] = tensors["state_slot_mapping_full"][0] full["position_ids"] = tensors["position_ids_full"][0] + full["query_start_loc"] = tensors["query_start_loc"][0] + full["local_request_ids"] = torch.cat( + [tensors["local_request_ids"][rank] for rank in range(tp_size)] + ) full["x_out"] = torch.zeros(token_count, HC_MULT, D, dtype=torch.float32) golden_prefill_attention_hca(full) @@ -1391,8 +1617,12 @@ def golden_prefill_attention_hca_cp(tensors): parser.add_argument("--compile-only", action="store_true", default=False) parser.add_argument("--start-pos", type=int, default=START_POS) parser.add_argument( - "--token-count", "--num-tokens", dest="token_count", type=int, default=PREFILL_SEQ, - help="Physical query-token extent across the group; must divide by --tp.", + "--token-count", "--num-tokens", dest="token_count", type=int, default=None, + help=f"B1 physical query-token extent across the group; defaults to {PREFILL_SEQ}. ragged2 is fixed at 8.", + ) + parser.add_argument( + "--case", choices=["b1", "ragged2"], default="b1", + help="Fixture case; ragged2 is the fixed two-request TP2 boundary case.", ) parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) parser.add_argument("--enable-dep-gen", action="store_true", default=False) @@ -1404,7 +1634,15 @@ def golden_prefill_attention_hca_cp(tensors): device_ids = [int(device) for device in args.device.split(",")] if len(device_ids) != TP_SIZE: parser.error(f"need exactly {TP_SIZE} devices, got {device_ids}") - if args.token_count % TP_SIZE != 0: + if args.case == "ragged2" and TP_SIZE != 2: + parser.error("--case ragged2 requires --tp 2") + if args.case == "ragged2" and args.start_pos != 0: + parser.error("--case ragged2 has fixed request starts and requires --start-pos 0") + if args.token_count is None: + args.token_count = 8 if args.case == "ragged2" else PREFILL_SEQ + if args.case == "ragged2" and args.token_count != 8: + parser.error("--case ragged2 has a fixed physical extent and requires --token-count 8") + if args.case == "b1" and args.token_count % TP_SIZE != 0: parser.error(f"--token-count must be a multiple of --tp={TP_SIZE}, got {args.token_count}") if TP_SIZE == 1: @@ -1432,9 +1670,14 @@ def golden_prefill_attention_hca_cp(tensors): else: from pypto.ir.distributed_compiled_program import DistributedConfig + specs = ( + build_ragged2_cp_tensor_specs(TP_SIZE) + if args.case == "ragged2" + else build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE) + ) result = run_jit( fn=l3_prefill_attention_hca_cp, - specs=build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE), + specs=specs, golden_fn=golden_prefill_attention_hca_cp, compile_cfg=dict( dump_passes=args.dump_passes, diff --git a/models/deepseek_v4_flash_dspark/prefill_indexer.py b/models/deepseek_v4_flash_dspark/prefill_indexer.py index 2403d9858..02190817e 100644 --- a/models/deepseek_v4_flash_dspark/prefill_indexer.py +++ b/models/deepseek_v4_flash_dspark/prefill_indexer.py @@ -29,6 +29,7 @@ golden_prefill_indexer_compressor, prefill_indexer_compressor, ) +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN PREFILL_MAX_TOKENS = 8192 @@ -163,6 +164,7 @@ def _topk_leaf( @pl.jit.incore def _topk_group_wave( position_ids: pl.Tensor, + local_request_ids: pl.Tensor, score_arena: pl.Tensor, pair_arena: pl.Tensor, tile_base: pl.Scalar[pl.INDEX], @@ -173,7 +175,10 @@ def _topk_group_wave( global_group_base = 0 for query in pl.range(tile_rows): position = pl.read(position_ids, [tile_base + query]) - visible_count = pl.max(pl.min((position + 1) // COMPRESS_RATIO, INDEXER_MAX_CANDIDATES), 0) + request_id = pl.read(local_request_ids, [tile_base + query]) + visible_count = 0 + if request_id >= 0: + visible_count = pl.max(pl.min((position + 1) // COMPRESS_RATIO, INDEXER_MAX_CANDIDATES), 0) leaf_count = (visible_count + TOPK_LEAF_TILE - 1) // TOPK_LEAF_TILE group_count = (leaf_count + TOPK_GROUP_TILE - 1) // TOPK_GROUP_TILE base_mod = global_group_base % TOPK_GROUP_WORKERS @@ -208,6 +213,7 @@ def _topk_group_wave( @pl.jit.incore def _topk_query_merge( position_ids: pl.Tensor, + local_request_ids: pl.Tensor, pair_arena: pl.Tensor, topk_indices: pl.Tensor, tile_base: pl.Scalar[pl.INDEX], @@ -216,7 +222,10 @@ def _topk_query_merge( query = pl.tile.get_block_idx() output_query = tile_base + query position = pl.read(position_ids, [output_query]) - visible_count = pl.max(pl.min((position + 1) // COMPRESS_RATIO, INDEXER_MAX_CANDIDATES), 0) + request_id = pl.read(local_request_ids, [output_query]) + visible_count = 0 + if request_id >= 0: + visible_count = pl.max(pl.min((position + 1) // COMPRESS_RATIO, INDEXER_MAX_CANDIDATES), 0) empty_indices = pl.tile.full([1, IDX_TOPK], dtype=pl.INT32, value=-1) pl.store(empty_indices, [output_query, 0], topk_indices) @@ -254,6 +263,7 @@ def _prefill_indexer_score_topk( idx_kv_cache: pl.Tensor, idx_kv_scale: pl.Tensor, idx_block_table: pl.Tensor, + local_request_ids: pl.Tensor, position_ids: pl.Tensor, topk_indices: pl.Tensor, score_arena: pl.Tensor, @@ -276,7 +286,10 @@ def _prefill_indexer_score_topk( for query in pl.range(tile_rows): output_query = tile_base + query position = pl.read(position_ids, [output_query]) - visible_count = pl.max(pl.min((position + 1) // COMPRESS_RATIO, INDEXER_MAX_CANDIDATES), 0) + request_id = pl.read(local_request_ids, [output_query]) + visible_count = 0 + if request_id >= 0: + visible_count = pl.max(pl.min((position + 1) // COMPRESS_RATIO, INDEXER_MAX_CANDIDATES), 0) leaf_count = (visible_count + TOPK_LEAF_TILE - 1) // TOPK_LEAF_TILE base_mod = global_leaf_base % TOPK_SCORE_WORKERS first_leaf = (worker + base_mod) % TOPK_SCORE_WORKERS @@ -292,7 +305,9 @@ def _prefill_indexer_score_topk( page_begin = page * BLOCK_SIZE logical_row = logical_begin + page_begin logical_page = logical_row // BLOCK_SIZE - physical_block_raw = pl.read(idx_block_table, [logical_page]) + physical_block_raw = pl.cast(-1, pl.INT32) + if request_id >= 0: + physical_block_raw = pl.read(idx_block_table, [request_id, logical_page]) score_valid = pl.full([1, BLOCK_SIZE], dtype=pl.FP32, value=FP32_NEG_INF) if physical_block_raw >= 0 and physical_block_raw < idx_block_num: physical_block = pl.cast(physical_block_raw, pl.INDEX) @@ -316,10 +331,10 @@ def _prefill_indexer_score_topk( global_leaf_base = global_leaf_base + leaf_count with pl.spmd(TOPK_GROUP_WORKERS, name_hint="prefill_idx_topk_group_wave", deps=[score_tid]) as topk_tid: - _topk_group_wave(position_ids, score_arena, pair_arena, tile_base, tile_rows) + _topk_group_wave(position_ids, local_request_ids, score_arena, pair_arena, tile_base, tile_rows) with pl.spmd(tile_rows, name_hint="prefill_idx_topk_query_merge", deps=[topk_tid]) as merge_tid: - _topk_query_merge(position_ids, pair_arena, topk_indices, tile_base) + _topk_query_merge(position_ids, local_request_ids, pair_arena, topk_indices, tile_base) completion[0] = merge_tid return topk_indices @@ -338,7 +353,8 @@ def _prefill_indexer_dense_tile( hadamard: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_kv_cache: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8], idx_kv_scale: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_topk_indices: pl.Out[pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32]], position_ids: pl.Tensor[[T_DYN], pl.INT32], rope_dup_idx_template: pl.Tensor[[1, ROPE_HEAD_DIM], pl.INT32], @@ -603,6 +619,7 @@ def _prefill_indexer_dense_tile( _prefill_indexer_score_topk( qr_hadamard_i8, qr_hadamard_scale_dq, weights, idx_kv_cache, idx_kv_scale, idx_block_table, + local_request_ids, position_ids, cmp_topk_indices, score_arena, pair_arena, selection_completion, @@ -613,6 +630,7 @@ def _prefill_indexer_dense_tile( @pl.jit.inline(auto_scope=False) def prefill_indexer( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], qr: pl.Tensor[[T_DYN, Q_LORA], pl.INT8], qr_scale: pl.Tensor[[T_DYN, 1], pl.FP32], wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], @@ -626,23 +644,25 @@ def prefill_indexer( inner_compress_state: pl.InOut[ pl.Tensor[[INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32] ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], inner_wkv: pl.Tensor[[INNER_OUT_DIM, D], pl.BF16], inner_wgate: pl.Tensor[[INNER_OUT_DIM, D], pl.BF16], inner_ape: pl.Tensor[[COMPRESS_RATIO, INNER_OUT_DIM], pl.FP32], inner_norm_w: pl.Tensor[[INNER_HEAD_DIM], pl.BF16], idx_kv_cache: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32]], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], cmp_topk_indices: pl.Out[pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32]], position_ids: pl.Tensor[[T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], idx_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], inner_state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], ): - """Compress and score one contiguous position-monotonic token run.""" + """Compress and score one packed ragged prefill stream.""" compressor_completion = pl.array.create(1, pl.TASK_ID) prefill_indexer_compressor( x, + query_start_loc, inner_compress_state, inner_compress_state_block_table, inner_wkv, inner_wgate, inner_ape, inner_norm_w, cmp_freqs_cos, cmp_freqs_sin, @@ -659,7 +679,7 @@ def prefill_indexer( hadamard, idx_kv_cache, idx_kv_scale, idx_block_table, cmp_topk_indices, - position_ids, + position_ids, local_request_ids, compressor_completion, ) return idx_kv_cache, idx_kv_scale, cmp_topk_indices @@ -678,9 +698,10 @@ def prefill_indexer_query( hadamard: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], idx_kv_cache: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8], idx_kv_scale: pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], cmp_topk_indices: pl.Out[pl.Tensor[[Q_T_DYN, IDX_TOPK], pl.INT32]], position_ids: pl.Tensor[[Q_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[Q_T_DYN], pl.INT32], completion: pl.Array[1, pl.TASK_ID], ): """Score local queries against the published indexer cache.""" @@ -717,6 +738,7 @@ def prefill_indexer_query( cos, sin, hadamard, idx_kv_cache, idx_kv_scale, idx_block_table, + local_request_ids, cmp_topk_indices, position_ids, rope_dup_idx_template, rope_swap_idx_template, rope_sign_template, score_arena, pair_arena, selection_completion, @@ -770,32 +792,33 @@ def golden_prefill_indexer_core(tensors): import torch token_count = int(tensors["x"].shape[0]) - compressor_tensors = { - "x": tensors["x"], - "kv": torch.zeros( - max(1, (token_count + COMPRESS_RATIO - 1) // COMPRESS_RATIO), - IDX_HEAD_DIM, - dtype=torch.int8, - ), - "compress_state": tensors["inner_compress_state"], - "inner_compress_state_block_table": tensors["inner_compress_state_block_table"], - "wkv": tensors["inner_wkv"], - "wgate": tensors["inner_wgate"], - "ape": tensors["inner_ape"], - "norm_w": tensors["inner_norm_w"], - "cmp_freqs_cos": tensors["cmp_freqs_cos"], - "cmp_freqs_sin": tensors["cmp_freqs_sin"], - "hadamard": tensors["hadamard"], - "idx_kv_cache": tensors["idx_kv_cache"], - "idx_kv_scale": tensors["idx_kv_scale"], - "idx_block_table": tensors["idx_block_table"], - "position_ids": tensors["position_ids"], - "idx_slot_mapping": tensors["idx_slot_mapping"], - "inner_state_slot_mapping": tensors["inner_state_slot_mapping"], - } - golden_prefill_indexer_compressor(compressor_tensors) - tensors["idx_kv_cache"][:] = compressor_tensors["idx_kv_cache"] - tensors["idx_kv_scale"][:] = compressor_tensors["idx_kv_scale"] + query_start_loc = tensors["query_start_loc"] + for request in range(query_start_loc.numel() - 1): + request_start = int(query_start_loc[request].item()) + request_end = int(query_start_loc[request + 1].item()) + if request_end <= request_start: + continue + request_rows = slice(request_start, request_end) + golden_prefill_indexer_compressor( + { + "x": tensors["x"][request_rows], + "compress_state": tensors["inner_compress_state"], + "inner_compress_state_block_table": tensors["inner_compress_state_block_table"][request : request + 1], + "wkv": tensors["inner_wkv"], + "wgate": tensors["inner_wgate"], + "ape": tensors["inner_ape"], + "norm_w": tensors["inner_norm_w"], + "cmp_freqs_cos": tensors["cmp_freqs_cos"][request_rows], + "cmp_freqs_sin": tensors["cmp_freqs_sin"][request_rows], + "hadamard": tensors["hadamard"], + "idx_kv_cache": tensors["idx_kv_cache"], + "idx_kv_scale": tensors["idx_kv_scale"], + "idx_block_table": tensors["idx_block_table"][request : request + 1], + "position_ids": tensors["position_ids"][request_rows], + "idx_slot_mapping": tensors["idx_slot_mapping"][request_rows], + "inner_state_slot_mapping": tensors["inner_state_slot_mapping"][request_rows], + } + ) # Lightning-indexer scores with per-token causal top-k. position_ids = tensors["position_ids"].long() @@ -815,13 +838,7 @@ def golden_prefill_indexer_core(tensors): cache_flat_i8 = tensors["idx_kv_cache"].reshape(-1, IDX_HEAD_DIM) scale_flat = tensors["idx_kv_scale"].float().reshape(-1, 1) idx_block_table = tensors["idx_block_table"] - logical_rows = torch.arange(max_visible, dtype=torch.int64) - physical_pages = idx_block_table[logical_rows // BLOCK_SIZE].to(torch.int64) - valid_pages = (physical_pages >= 0) & (physical_pages < tensors["idx_kv_cache"].shape[0]) - safe_pages = physical_pages.clamp(min=0, max=tensors["idx_kv_cache"].shape[0] - 1) - physical_rows = safe_pages * BLOCK_SIZE + logical_rows % BLOCK_SIZE - kv_i8 = cache_flat_i8[physical_rows] - kv_sc = scale_flat[physical_rows, 0] + local_request_ids = tensors["local_request_ids"] # Query tiles and 8192-candidate score leaves. for tile_base in range(0, token_count, PREFILL_DENSE_TILE): @@ -852,9 +869,19 @@ def golden_prefill_indexer_core(tensors): q_sc = q_sc.view(tile_rows, IDX_N_HEADS, 1) for local_t in range(tile_rows): global_t = tile_base + local_t + request_id = int(local_request_ids[global_t].item()) + if request_id < 0: + continue visible_t = int(visible[global_t].item()) if visible_t <= 0: continue + logical_rows = torch.arange(visible_t, dtype=torch.int64) + physical_pages = idx_block_table[request_id, logical_rows // BLOCK_SIZE].to(torch.int64) + valid_pages = (physical_pages >= 0) & (physical_pages < tensors["idx_kv_cache"].shape[0]) + safe_pages = physical_pages.clamp(min=0, max=tensors["idx_kv_cache"].shape[0] - 1) + physical_rows = safe_pages * BLOCK_SIZE + logical_rows % BLOCK_SIZE + kv_i8 = cache_flat_i8[physical_rows] + kv_sc = scale_flat[physical_rows, 0] running_scores = torch.empty(0, dtype=torch.float32) running_indices = torch.empty(0, dtype=torch.int64) for begin in range(0, visible_t, 8192): @@ -885,6 +912,7 @@ def golden_prefill_indexer(tensors): @pl.jit def prefill_indexer_test( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], qr: pl.Tensor[[T_DYN, Q_LORA], pl.INT8], qr_scale: pl.Tensor[[T_DYN, 1], pl.FP32], wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], @@ -898,20 +926,22 @@ def prefill_indexer_test( inner_compress_state: pl.InOut[ pl.Tensor[[INNER_STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32] ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], inner_wkv: pl.Tensor[[INNER_OUT_DIM, D], pl.BF16], inner_wgate: pl.Tensor[[INNER_OUT_DIM, D], pl.BF16], inner_ape: pl.Tensor[[COMPRESS_RATIO, INNER_OUT_DIM], pl.FP32], inner_norm_w: pl.Tensor[[INNER_HEAD_DIM], pl.BF16], idx_kv_cache: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[IDX_BLOCK_NUM_DYN, BLOCK_SIZE, 1, 1], pl.FP32]], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], topk_idxs: pl.Out[pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32]], position_ids: pl.Tensor[[T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], idx_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], inner_state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], ): x.bind_dynamic(0, T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) qr.bind_dynamic(0, T_DYN) qr_scale.bind_dynamic(0, T_DYN) cos.bind_dynamic(0, T_DYN) @@ -919,15 +949,18 @@ def prefill_indexer_test( cmp_freqs_cos.bind_dynamic(0, T_DYN) cmp_freqs_sin.bind_dynamic(0, T_DYN) inner_compress_state.bind_dynamic(0, INNER_STATE_BLOCK_NUM_DYN) + inner_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) idx_kv_cache.bind_dynamic(0, IDX_BLOCK_NUM_DYN) idx_kv_scale.bind_dynamic(0, IDX_BLOCK_NUM_DYN) topk_idxs.bind_dynamic(0, T_DYN) position_ids.bind_dynamic(0, T_DYN) + local_request_ids.bind_dynamic(0, T_DYN) idx_slot_mapping.bind_dynamic(0, T_DYN) inner_state_slot_mapping.bind_dynamic(0, T_DYN) prefill_indexer( x, + query_start_loc, qr, qr_scale, wq_b, @@ -949,6 +982,7 @@ def prefill_indexer_test( idx_block_table, topk_idxs, position_ids, + local_request_ids, idx_slot_mapping, inner_state_slot_mapping, ) @@ -1007,7 +1041,7 @@ def build_tensor_specs(start_pos: int = START_POS, token_count: int = PREFILL_SE def init_inner_compress_state_block_table(): blocks = torch.arange(INNER_STATE_MAX_BLOCKS, dtype=torch.int64) - return ((blocks * 17 + 3) % CSA_INNER_STATE_PHYSICAL_BLOCKS).to(torch.int32) + return ((blocks * 17 + 3) % CSA_INNER_STATE_PHYSICAL_BLOCKS).to(torch.int32).unsqueeze(0) def state_row(abs_pos): if abs_pos < 0 or abs_pos >= MAX_SEQ_LEN: @@ -1059,7 +1093,7 @@ def _build_idx_hist(): c_flat = cache_i8.view(IDX_CACHE_BLOCK_NUM * BLOCK_SIZE, IDX_HEAD_DIM) s_flat = scale.view(IDX_CACHE_BLOCK_NUM * BLOCK_SIZE, 1) completed = start_pos // COMPRESS_RATIO - table = init_idx_block_table().to(torch.int64) + table = init_idx_block_table()[0].to(torch.int64) if completed > table.numel() * BLOCK_SIZE: raise ValueError("fixture historical compressed slots exceed the standalone idx block table") history_chunk = 16 * 1024 @@ -1089,7 +1123,10 @@ def init_idx_kv_scale(): return _idx_hist["scale"].clone() def init_idx_block_table(): - return torch.arange(IDX_CACHE_MAX_BLOCKS, dtype=torch.int32) + return torch.arange(IDX_CACHE_MAX_BLOCKS, dtype=torch.int32).unsqueeze(0) + + def init_local_request_ids(): + return torch.zeros(token_count, dtype=torch.int32) def init_position_ids(): return torch.arange(start_pos, start_pos + token_count, dtype=torch.int32) @@ -1119,7 +1156,7 @@ def init_idx_slot_mapping(): write_mask = (positions + 1) % COMPRESS_RATIO == 0 if write_mask.any(): compressed_slots = (positions[write_mask] + 1) // COMPRESS_RATIO - 1 - table = init_idx_block_table().to(torch.int64) + table = init_idx_block_table()[0].to(torch.int64) physical_pages = table[compressed_slots // BLOCK_SIZE] rows = physical_pages * BLOCK_SIZE + compressed_slots % BLOCK_SIZE if (physical_pages < 0).any(): @@ -1163,6 +1200,7 @@ def init_sin(): return [ TensorSpec("x", [token_count, D], torch.bfloat16, init_value=init_x), + TensorSpec("query_start_loc", [2], torch.int32, init_value=torch.tensor([0, token_count], dtype=torch.int32)), TensorSpec("qr", [token_count, Q_LORA], torch.int8, init_value=lambda: qr_i8), TensorSpec("qr_scale", [token_count, 1], torch.float32, init_value=lambda: qr_scale), TensorSpec("wq_b", [Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], torch.int8, init_value=lambda: wq_b_i8), @@ -1182,7 +1220,7 @@ def init_sin(): ), TensorSpec( "inner_compress_state_block_table", - [INNER_STATE_MAX_BLOCKS], + [1, INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table, ), @@ -1204,9 +1242,10 @@ def init_sin(): init_value=init_idx_kv_scale, is_output=True, ), - TensorSpec("idx_block_table", [IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), + TensorSpec("idx_block_table", [1, IDX_CACHE_MAX_BLOCKS], torch.int32, init_value=init_idx_block_table), TensorSpec("topk_idxs", [token_count, IDX_TOPK], torch.int32, is_output=True), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), + TensorSpec("local_request_ids", [token_count], torch.int32, init_value=init_local_request_ids), TensorSpec( "idx_slot_mapping", [token_count], @@ -1278,7 +1317,8 @@ def score_selected_indices(token_id, indices, expected_outputs, inputs): weights = (inputs["x"][token_id].float() @ inputs["weights_proj"].float()) * WEIGHTS_SCALE logical_rows = indices.to(torch.int64) - physical_pages = inputs["idx_block_table"][logical_rows // BLOCK_SIZE].to(torch.int64) + request_id = int(inputs["local_request_ids"][token_id].item()) + physical_pages = inputs["idx_block_table"][request_id, logical_rows // BLOCK_SIZE].to(torch.int64) physical_rows = physical_pages * BLOCK_SIZE + logical_rows % BLOCK_SIZE cache = expected_outputs["idx_kv_cache"].reshape(-1, IDX_HEAD_DIM) scales = expected_outputs["idx_kv_scale"].float().reshape(-1, 1) diff --git a/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py b/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py index a508b449e..6d007a926 100644 --- a/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py +++ b/models/deepseek_v4_flash_dspark/prefill_indexer_compressor.py @@ -20,8 +20,7 @@ INT8_SCALE_MAX, PREFILL_SEQ, ) - - +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN # Bounded physical-row tile for index projection/state updates. PREFILL_STATE_TILE = 512 @@ -243,41 +242,39 @@ def _prefill_indexer_compressor_tile( write_dst_map[0:1, 0:MAX_CMP_WRITES] = write_dst_tile write_src_map[0:1, 0:MAX_CMP_WRITES] = write_src_tile - # Scatter all physical rows before pooling. The fence read registers the - # previous tile's commit as this tile's predecessor in TensorMap. - with pl.spmd(tile_rows, name_hint="prefill_idx_c4_state_scatter_pre") as _scatter_tid: + # Carry the previous tile's commit into state scatter. + with pl.spmd(tile_rows, name_hint="prefill_idx_c4_state_scatter_pre"): scatter_local = pl.tile.get_block_idx() - scatter_order = pl.read(state_order_fence, [0]) - if scatter_order >= 0: - scatter_global = tile_base + scatter_local - state_row_raw = pl.read(inner_state_slot_mapping, [scatter_global]) - if state_row_raw >= 0: - state_row = pl.cast(state_row_raw, pl.INDEX) - scatter_pos = pl.read(position_ids, [scatter_global]) - ape_slot = pl.cast(scatter_pos % COMPRESS_RATIO, pl.INDEX) - for scatter_ob in pl.range(OUT_DIM // OUT_TILE): - scatter_o0 = scatter_ob * OUT_TILE - ape_row = ape[ - ape_slot : ape_slot + 1, - scatter_o0 : scatter_o0 + OUT_TILE, - ] - compress_state_flat[ - state_row : state_row + 1, - scatter_o0 : scatter_o0 + OUT_TILE, - ] = kv_proj_scratch[ + _state_order_anchor = pl.read(state_order_fence, [0]) + scatter_global = tile_base + scatter_local + state_row_raw = pl.read(inner_state_slot_mapping, [scatter_global]) + if state_row_raw >= 0: + state_row = pl.cast(state_row_raw, pl.INDEX) + scatter_pos = pl.read(position_ids, [scatter_global]) + ape_slot = pl.cast(scatter_pos % COMPRESS_RATIO, pl.INDEX) + for scatter_ob in pl.range(OUT_DIM // OUT_TILE): + scatter_o0 = scatter_ob * OUT_TILE + ape_row = ape[ + ape_slot : ape_slot + 1, + scatter_o0 : scatter_o0 + OUT_TILE, + ] + compress_state_flat[ + state_row : state_row + 1, + scatter_o0 : scatter_o0 + OUT_TILE, + ] = kv_proj_scratch[ + scatter_local : scatter_local + 1, + scatter_o0 : scatter_o0 + OUT_TILE, + ] + compress_state_flat[ + state_row : state_row + 1, + OUT_DIM + scatter_o0 : OUT_DIM + scatter_o0 + OUT_TILE, + ] = pl.add( + score_proj_scratch[ scatter_local : scatter_local + 1, scatter_o0 : scatter_o0 + OUT_TILE, - ] - compress_state_flat[ - state_row : state_row + 1, - OUT_DIM + scatter_o0 : OUT_DIM + scatter_o0 + OUT_TILE, - ] = pl.add( - score_proj_scratch[ - scatter_local : scatter_local + 1, - scatter_o0 : scatter_o0 + OUT_TILE, - ], - ape_row, - ) + ], + ape_row, + ) for pool_idx in pl.spmd( MAX_CMP_WRITES * (HEAD_DIM // HEAD_D_TILE), @@ -568,10 +565,11 @@ def _prefill_indexer_compressor_tile( @pl.jit.inline(auto_scope=False) def prefill_indexer_compressor( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32] ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], wkv: pl.Tensor[[OUT_DIM, D], pl.BF16], wgate: pl.Tensor[[OUT_DIM, D], pl.BF16], ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32], @@ -586,8 +584,8 @@ def prefill_indexer_compressor( inner_state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], completion: pl.Array[1, pl.TASK_ID], ): - """Run one physical request through ordered 512-row compressor tiles.""" - t_dim = pl.tensor.dim(x, 0) + """Compress packed requests independently through ordered 512-row state tiles.""" + request_count = pl.tensor.dim(query_start_loc, 0) - 1 rope_dup_idx_template = pl.create_tensor([PACKED_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.INT32) rope_swap_idx_template = pl.create_tensor([PACKED_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.INT32) rope_sign_template = pl.create_tensor([PACKED_RMS_TILE, ROPE_HEAD_DIM], dtype=pl.FP32) @@ -615,35 +613,43 @@ def prefill_indexer_compressor( state_order_fence = pl.create_tensor([1], dtype=pl.INT32) with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_idx_c4_state_order_init"): pl.write(state_order_fence, [0], pl.cast(0, pl.INT32)) - for tile_base in pl.range(0, t_dim, PREFILL_STATE_TILE): - tile_rows = pl.min(PREFILL_STATE_TILE, t_dim - tile_base) - with pl.scope(): - _prefill_indexer_compressor_tile( - x, - compress_state, - inner_compress_state_block_table, - wkv, - wgate, - ape, - norm_w, - cmp_freqs_cos, - cmp_freqs_sin, - hadamard, - idx_kv_cache, - idx_kv_scale, - position_ids, - idx_slot_mapping, - inner_state_slot_mapping, - rope_dup_idx_template, - rope_swap_idx_template, - rope_sign_template, - state_order_fence, - tile_base, - tile_rows, - ) + for request in pl.range(request_count): + request_start = pl.cast(pl.read(query_start_loc, [request]), pl.INDEX) + request_end = pl.cast(pl.read(query_start_loc, [request + 1]), pl.INDEX) + request_table = inner_compress_state_block_table[request] + for request_offset in pl.range(0, request_end - request_start, PREFILL_STATE_TILE): + tile_base = request_start + request_offset + tile_rows = pl.min(PREFILL_STATE_TILE, request_end - tile_base) + with pl.scope(): + _prefill_indexer_compressor_tile( + x, + compress_state, + request_table, + wkv, + wgate, + ape, + norm_w, + cmp_freqs_cos, + cmp_freqs_sin, + hadamard, + idx_kv_cache, + idx_kv_scale, + position_ids, + idx_slot_mapping, + inner_state_slot_mapping, + rope_dup_idx_template, + rope_swap_idx_template, + rope_sign_template, + state_order_fence, + tile_base, + tile_rows, + ) # Compressor completion fence. - with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_idx_c4_complete") as completion_tid: + with pl.at( + level=pl.Level.CORE_GROUP, + name_hint="prefill_idx_c4_complete", + ) as completion_tid: fence_sample = pl.read(state_order_fence, [0]) completion_bit = pl.cast(fence_sample == fence_sample, pl.INT32) pl.write(state_order_fence, [0], completion_bit * completion_bit) @@ -654,10 +660,11 @@ def prefill_indexer_compressor( @pl.jit def prefill_indexer_compressor_test( x: pl.Tensor[[T_DYN, D], pl.BF16], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], compress_state: pl.InOut[ pl.Tensor[[STATE_BLOCK_NUM_DYN, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], pl.FP32] ], - inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], wkv: pl.Tensor[[OUT_DIM, D], pl.BF16], wgate: pl.Tensor[[OUT_DIM, D], pl.BF16], ape: pl.Tensor[[COMPRESS_RATIO, OUT_DIM], pl.FP32], @@ -672,7 +679,9 @@ def prefill_indexer_compressor_test( inner_state_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], ): x.bind_dynamic(0, T_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) compress_state.bind_dynamic(0, STATE_BLOCK_NUM_DYN) + inner_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) idx_kv_cache.bind_dynamic(0, IDX_BLOCK_NUM_DYN) idx_kv_scale.bind_dynamic(0, IDX_BLOCK_NUM_DYN) cmp_freqs_cos.bind_dynamic(0, T_DYN) @@ -684,6 +693,7 @@ def prefill_indexer_compressor_test( completion = pl.array.create(1, pl.TASK_ID) return prefill_indexer_compressor( x, + query_start_loc, compress_state, inner_compress_state_block_table, wkv, @@ -714,7 +724,7 @@ def golden_prefill_indexer_compressor(tensors): ) kv_state_flat = compress_state_flat[:, :OUT_DIM] score_state_flat = compress_state_flat[:, OUT_DIM:] - state_block_table = tensors["inner_compress_state_block_table"] + state_block_table = tensors["inner_compress_state_block_table"][0] idx_kv_cache = tensors["idx_kv_cache"] # C8: INT8 KV idx_kv_scale = tensors["idx_kv_scale"] # C8: per-position FP32 dequant scale cache_rows = idx_kv_cache.view(idx_kv_cache.shape[0] * BLOCK_SIZE, 1, HEAD_DIM)[:, 0, :] @@ -828,7 +838,7 @@ def build_tensor_specs(start_pos: int = START_POS, token_count: int = PREFILL_SE def init_inner_compress_state_block_table(): logical_blocks = torch.arange(INNER_STATE_MAX_BLOCKS, dtype=torch.int64) - return ((logical_blocks * 17 + 3) % CSA_INNER_STATE_PHYSICAL_BLOCKS).to(torch.int32) + return ((logical_blocks * 17 + 3) % CSA_INNER_STATE_PHYSICAL_BLOCKS).to(torch.int32).unsqueeze(0) def state_row(abs_pos): if abs_pos < 0 or abs_pos >= MAX_SEQ_LEN: @@ -923,6 +933,7 @@ def init_inner_state_slot_mapping(): return [ TensorSpec("x", [token_count, D], torch.bfloat16, init_value=init_x), + TensorSpec("query_start_loc", [2], torch.int32, init_value=torch.tensor([0, token_count], dtype=torch.int32)), TensorSpec( "compress_state", [INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, COMPRESS_STATE_DIM], @@ -932,7 +943,7 @@ def init_inner_state_slot_mapping(): ), TensorSpec( "inner_compress_state_block_table", - [INNER_STATE_MAX_BLOCKS], + [1, INNER_STATE_MAX_BLOCKS], torch.int32, init_value=init_inner_compress_state_block_table, ), diff --git a/models/deepseek_v4_flash_dspark/prefill_layer.py b/models/deepseek_v4_flash_dspark/prefill_layer.py index 56789af0e..3d55aadd3 100644 --- a/models/deepseek_v4_flash_dspark/prefill_layer.py +++ b/models/deepseek_v4_flash_dspark/prefill_layer.py @@ -100,6 +100,7 @@ prefill_attention_hca_cp, ) from prefill_cp_token_allgather import PREFILL_GROUP_CAP, TP_SIZE +from prefill_metadata import QUERY_START_LOC_DYN, REQUESTS_DYN, lower_local_request_ids from prefill_swa import golden_prefill_attention_swa, prefill_attention_swa_cp @@ -116,6 +117,7 @@ @pl.jit(auto_scope=False) def prefill_layer_attention( x_hc: pl.Tensor[[FWD_GROUP_TOKENS_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[3], pl.FP32], hc_attn_base: pl.Tensor[[MIX_HC], pl.FP32], @@ -139,13 +141,13 @@ def prefill_layer_attention( hca_cmp_ape: pl.Tensor[[128, HCA_MAIN_OUT_DIM], pl.FP32], hca_cmp_norm_w: pl.Tensor[[HEAD_DIM], pl.BF16], hca_compress_state: pl.InOut[pl.Tensor[[HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, HCA_COMPRESS_STATE_DIM], pl.FP32]], - hca_compress_state_block_table: pl.Tensor[[HCA_STATE_MAX_BLOCKS], pl.INT32], + hca_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], csa_cmp_wkv: pl.Tensor[[CSA_MAIN_OUT_DIM, D], pl.BF16], csa_cmp_wgate: pl.Tensor[[CSA_MAIN_OUT_DIM, D], pl.BF16], csa_cmp_ape: pl.Tensor[[4, CSA_MAIN_OUT_DIM], pl.FP32], csa_cmp_norm_w: pl.Tensor[[HEAD_DIM], pl.BF16], csa_compress_state: pl.InOut[pl.Tensor[[CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, CSA_COMPRESS_STATE_DIM], pl.FP32]], - csa_compress_state_block_table: pl.Tensor[[CSA_STATE_MAX_BLOCKS], pl.INT32], + csa_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], csa_hadamard_idx: pl.Tensor[[IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], csa_idx_wq_b: pl.Tensor[[Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], csa_idx_wq_b_scale: pl.Tensor[[IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -155,17 +157,17 @@ def prefill_layer_attention( csa_inner_ape: pl.Tensor[[4, INNER_OUT_DIM], pl.FP32], csa_inner_norm_w: pl.Tensor[[IDX_HEAD_DIM], pl.BF16], csa_inner_compress_state: pl.InOut[pl.Tensor[[INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32]], - csa_inner_compress_state_block_table: pl.Tensor[[INNER_STATE_MAX_BLOCKS], pl.INT32], + csa_inner_compress_state_block_table: pl.Tensor[[REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[CSA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - ori_block_table: pl.Tensor[[SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[FWD_GROUP_TOKENS_DYN], pl.INT64], hca_cmp_kv: pl.InOut[pl.Tensor[[HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], csa_cmp_kv: pl.InOut[pl.Tensor[[CSA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - hca_cmp_block_table: pl.Tensor[[HCA_CMP_MAX_BLOCKS], pl.INT32], - csa_cmp_block_table: pl.Tensor[[CSA_CMP_MAX_BLOCKS], pl.INT32], + hca_cmp_block_table: pl.Tensor[[REQUESTS_DYN, HCA_CMP_MAX_BLOCKS], pl.INT32], + csa_cmp_block_table: pl.Tensor[[REQUESTS_DYN, CSA_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.InOut[pl.Tensor[[IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], pl.FP32]], - idx_block_table: pl.Tensor[[IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[FWD_TOKENS_DYN], pl.INT32], position_ids_full: pl.Tensor[[FWD_GROUP_TOKENS_DYN], pl.INT32], hca_cmp_slot_mapping_full: pl.Tensor[[FWD_GROUP_TOKENS_DYN], pl.INT64], @@ -191,6 +193,14 @@ def prefill_layer_attention( ): """Run one selected DSA-CP attention kind.""" x_hc.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) + hca_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) + csa_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) + csa_inner_compress_state_block_table.bind_dynamic(0, REQUESTS_DYN) + ori_block_table.bind_dynamic(0, REQUESTS_DYN) + hca_cmp_block_table.bind_dynamic(0, REQUESTS_DYN) + csa_cmp_block_table.bind_dynamic(0, REQUESTS_DYN) + idx_block_table.bind_dynamic(0, REQUESTS_DYN) swa_freqs_cos.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) swa_freqs_sin.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) compressed_freqs_cos.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) @@ -209,6 +219,9 @@ def prefill_layer_attention( csa_state_slot_mapping_full.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) csa_inner_state_slot_mapping_full.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) attn_stage.bind_dynamic(0, FWD_GROUP_TOKENS_DYN) + local_tokens = pl.tensor.dim(position_ids_local, 0) + local_request_ids = pl.create_tensor([local_tokens], dtype=pl.INT32) + lower_local_request_ids(query_start_loc, local_request_ids, tp_rank * local_tokens) wo_a_full = pl.create_tensor([O_PROJ_SCRATCH_GROUPS, O_PROJ_SCRATCH_RANK, O_PROJ_SCRATCH_INPUT], dtype=pl.BF16) wo_b_full = pl.create_tensor([O_PROJ_SCRATCH_D, O_PROJ_SCRATCH_COLS], dtype=pl.INT8) o_proj_order_fence = pl.create_tensor([1], dtype=pl.INT32) @@ -225,7 +238,7 @@ def prefill_layer_attention( wkv, gamma_cq, gamma_ckv, swa_freqs_cos, swa_freqs_sin, kv_cache, ori_block_table, ori_slot_mapping_full, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, attn_sink, wo_a, wo_b, wo_b_scale, wo_a_full, wo_b_full, @@ -239,6 +252,7 @@ def prefill_layer_attention( elif layer_id == 2: attn_stage, gather_signal = prefill_attention_csa_cp( x_hc, + query_start_loc, hc_attn_fn, hc_attn_scale, hc_attn_base, attn_norm_w, wq_a, wq_b, wq_b_scale, @@ -256,7 +270,7 @@ def prefill_layer_attention( kv_cache, ori_block_table, ori_slot_mapping_full, csa_cmp_kv, csa_cmp_block_table, idx_kv_cache, idx_kv_scale, idx_block_table, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, csa_cmp_slot_mapping_full, csa_idx_slot_mapping_full, csa_state_slot_mapping_full, csa_inner_state_slot_mapping_full, attn_sink, @@ -272,6 +286,7 @@ def prefill_layer_attention( else: attn_stage, gather_signal = prefill_attention_hca_cp( x_hc, + query_start_loc, local_request_ids, hc_attn_fn, hc_attn_scale, hc_attn_base, attn_norm_w, wq_a, wq_b, wq_b_scale, @@ -407,6 +422,7 @@ def prefill_layer_moe( @pl.jit.host def l3_prefill_layer( x_hc: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN, HC_MULT, D], pl.FP32], + query_start_loc: pl.Tensor[[N_RANKS, QUERY_START_LOC_DYN], pl.INT32], hc_attn_fn: pl.Tensor[[N_RANKS, MIX_HC, HC_DIM], pl.FP32], hc_attn_scale: pl.Tensor[[N_RANKS, 3], pl.FP32], hc_attn_base: pl.Tensor[[N_RANKS, MIX_HC], pl.FP32], @@ -430,13 +446,13 @@ def l3_prefill_layer( hca_cmp_ape: pl.Tensor[[N_RANKS, 128, HCA_MAIN_OUT_DIM], pl.FP32], hca_cmp_norm_w: pl.Tensor[[N_RANKS, HEAD_DIM], pl.BF16], hca_compress_state: pl.InOut[pl.Tensor[[N_RANKS, HCA_STATE_BLOCK_NUM, HCA_STATE_BLOCK_SIZE, HCA_COMPRESS_STATE_DIM], pl.FP32]], - hca_compress_state_block_table: pl.Tensor[[N_RANKS, HCA_STATE_MAX_BLOCKS], pl.INT32], + hca_compress_state_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, HCA_STATE_MAX_BLOCKS], pl.INT32], csa_cmp_wkv: pl.Tensor[[N_RANKS, CSA_MAIN_OUT_DIM, D], pl.BF16], csa_cmp_wgate: pl.Tensor[[N_RANKS, CSA_MAIN_OUT_DIM, D], pl.BF16], csa_cmp_ape: pl.Tensor[[N_RANKS, 4, CSA_MAIN_OUT_DIM], pl.FP32], csa_cmp_norm_w: pl.Tensor[[N_RANKS, HEAD_DIM], pl.BF16], csa_compress_state: pl.InOut[pl.Tensor[[N_RANKS, CSA_STATE_BLOCK_NUM, CSA_STATE_BLOCK_SIZE, CSA_COMPRESS_STATE_DIM], pl.FP32]], - csa_compress_state_block_table: pl.Tensor[[N_RANKS, CSA_STATE_MAX_BLOCKS], pl.INT32], + csa_compress_state_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, CSA_STATE_MAX_BLOCKS], pl.INT32], csa_hadamard_idx: pl.Tensor[[N_RANKS, IDX_HEAD_DIM, IDX_HEAD_DIM], pl.BF16], csa_idx_wq_b: pl.Tensor[[N_RANKS, Q_LORA, IDX_N_HEADS * IDX_HEAD_DIM], pl.INT8], csa_idx_wq_b_scale: pl.Tensor[[N_RANKS, IDX_N_HEADS * IDX_HEAD_DIM], pl.FP32], @@ -446,17 +462,17 @@ def l3_prefill_layer( csa_inner_ape: pl.Tensor[[N_RANKS, 4, INNER_OUT_DIM], pl.FP32], csa_inner_norm_w: pl.Tensor[[N_RANKS, IDX_HEAD_DIM], pl.BF16], csa_inner_compress_state: pl.InOut[pl.Tensor[[N_RANKS, INNER_STATE_BLOCK_NUM, INNER_STATE_BLOCK_SIZE, INNER_COMPRESS_STATE_DIM], pl.FP32]], - csa_inner_compress_state_block_table: pl.Tensor[[N_RANKS, INNER_STATE_MAX_BLOCKS], pl.INT32], + csa_inner_compress_state_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, INNER_STATE_MAX_BLOCKS], pl.INT32], kv_cache: pl.InOut[pl.Tensor[[N_RANKS, CSA_ORI_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - ori_block_table: pl.Tensor[[N_RANKS, SPARSE_ORI_MAX_BLOCKS], pl.INT32], + ori_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, SPARSE_ORI_MAX_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN], pl.INT64], hca_cmp_kv: pl.InOut[pl.Tensor[[N_RANKS, HCA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], csa_cmp_kv: pl.InOut[pl.Tensor[[N_RANKS, CSA_CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - hca_cmp_block_table: pl.Tensor[[N_RANKS, HCA_CMP_MAX_BLOCKS], pl.INT32], - csa_cmp_block_table: pl.Tensor[[N_RANKS, CSA_CMP_MAX_BLOCKS], pl.INT32], + hca_cmp_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, HCA_CMP_MAX_BLOCKS], pl.INT32], + csa_cmp_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, CSA_CMP_MAX_BLOCKS], pl.INT32], idx_kv_cache: pl.InOut[pl.Tensor[[N_RANKS, IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, IDX_HEAD_DIM], pl.INT8]], idx_kv_scale: pl.InOut[pl.Tensor[[N_RANKS, IDX_CACHE_BLOCK_NUM, BLOCK_SIZE, 1, 1], pl.FP32]], - idx_block_table: pl.Tensor[[N_RANKS, IDX_CACHE_MAX_BLOCKS], pl.INT32], + idx_block_table: pl.Tensor[[N_RANKS, REQUESTS_DYN, IDX_CACHE_MAX_BLOCKS], pl.INT32], position_ids_local: pl.Tensor[[N_RANKS, FWD_TOKENS_DYN], pl.INT32], position_ids_full: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN], pl.INT32], hca_cmp_slot_mapping_full: pl.Tensor[[N_RANKS, FWD_GROUP_TOKENS_DYN], pl.INT64], @@ -499,6 +515,14 @@ def l3_prefill_layer( ): """Run one DSA-CP layer across all ranks.""" x_hc.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) + query_start_loc.bind_dynamic(1, QUERY_START_LOC_DYN) + hca_compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) + csa_compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) + csa_inner_compress_state_block_table.bind_dynamic(1, REQUESTS_DYN) + ori_block_table.bind_dynamic(1, REQUESTS_DYN) + hca_cmp_block_table.bind_dynamic(1, REQUESTS_DYN) + csa_cmp_block_table.bind_dynamic(1, REQUESTS_DYN) + idx_block_table.bind_dynamic(1, REQUESTS_DYN) swa_freqs_cos.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) swa_freqs_sin.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) compressed_freqs_cos.bind_dynamic(1, FWD_GROUP_TOKENS_DYN) @@ -555,6 +579,7 @@ def l3_prefill_layer( tp_rank = rank % TP_SIZE prefill_layer_attention( x_hc[rank], + query_start_loc[rank], hc_attn_fn[rank], hc_attn_scale[rank], hc_attn_base[rank], attn_norm_w[rank], wq_a[rank], wq_b[rank], wq_b_scale[rank], @@ -726,7 +751,11 @@ def init_full_x(): return torch.cat(groups, dim=0) specs_by_name = { - "x_hc": TensorSpec("x_hc", [N_RANKS, physical_tokens, HC_MULT, D], torch.float32, init_value=init_full_x) + "x_hc": TensorSpec("x_hc", [N_RANKS, physical_tokens, HC_MULT, D], torch.float32, init_value=init_full_x), + "query_start_loc": TensorSpec( + "query_start_loc", [N_RANKS, 2], torch.int32, + init_value=lambda: torch.tensor([0, token_count], dtype=torch.int32).repeat(N_RANKS, 1), + ), } padded_mapping_names = { "ori_slot_mapping_full", @@ -785,15 +814,29 @@ def init_attn_stage(): for name in _RESIDENT_WEIGHT_NAMES: specs_by_name[name].resident = "stacked" - tensor_order = (*[name for name in HOST_TENSOR_ORDER if name != "x_next"], *_STAGE_TENSOR_ORDER, "x_next") + tensor_order = ( + "x_hc", + *[name for name in HOST_TENSOR_ORDER if name not in {"x_hc", "x_next"}], + *_STAGE_TENSOR_ORDER, "x_next", + ) return [specs_by_name[name] for name in tensor_order] + [ScalarSpec("layer_id", torch.int32, layer_id)] def _attention_golden_tensors(tensors, rank, layer_id, x_out): + import torch + wo_a_full = tensors["wo_a"][rank : rank + TP_SIZE].reshape(O_GROUPS, O_LORA, O_GROUP_IN) wo_b_full = tensors["wo_b"][rank : rank + TP_SIZE].permute(1, 0, 2).reshape(D, O_GROUPS * O_LORA) + query_start_loc = tensors["query_start_loc"][rank] + local_request_ids = torch.full((tensors["x_hc"].shape[1],), -1, dtype=torch.int32) + for request_id in range(query_start_loc.numel() - 1): + request_start = int(query_start_loc[request_id]) + request_end = int(query_start_loc[request_id + 1]) + local_request_ids[request_start:request_end] = request_id common = { "x_hc": tensors["x_hc"][rank], + "query_start_loc": query_start_loc, + "local_request_ids": local_request_ids, "hc_attn_fn": tensors["hc_attn_fn"][rank], "hc_attn_scale": tensors["hc_attn_scale"][rank], "hc_attn_base": tensors["hc_attn_base"][rank], diff --git a/models/deepseek_v4_flash_dspark/prefill_metadata.py b/models/deepseek_v4_flash_dspark/prefill_metadata.py new file mode 100644 index 000000000..30f5e33e7 --- /dev/null +++ b/models/deepseek_v4_flash_dspark/prefill_metadata.py @@ -0,0 +1,116 @@ +# Copyright (c) PyPTO Contributors. +# This program is free software, you can redistribute it and/or modify it under the terms and conditions of +# CANN Open Software License Agreement Version 2.0 (the "License"). +# Please refer to the License for details. You may not use this file except in compliance with the License. +# THIS SOFTWARE IS PROVIDED ON AN "AS IS" BASIS, WITHOUT WARRANTIES OF ANY KIND, EITHER EXPRESS OR IMPLIED, +# INCLUDING BUT NOT LIMITED TO NON-INFRINGEMENT, MERCHANTABILITY, OR FITNESS FOR A PARTICULAR PURPOSE. +# See LICENSE in the root of the software repository for the full text of the License. +# ----------------------------------------------------------------------------------------------------------- +"""Shared dynamic request axes and device-side lowering for packed prefill.""" + +import pypto.language as pl + + +# Dynamic dimensions shared by packed prefill metadata consumers. +REQUESTS_DYN = pl.dynamic("PREFILL_METADATA_REQUESTS_DYN") +QUERY_START_LOC_DYN = pl.dynamic("PREFILL_METADATA_QUERY_START_LOC_DYN") +LOCAL_TOKENS_DYN = pl.dynamic("PREFILL_METADATA_LOCAL_TOKENS_DYN") + + +@pl.jit.inline +def lower_local_request_ids( + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], + local_request_ids: pl.Tensor[[LOCAL_TOKENS_DYN], pl.INT32], + local_base: pl.Scalar[pl.INT32], +): + """Lower packed query starts into request ids for one TP rank.""" + request_count = pl.tensor.dim(query_start_loc, 0) - 1 + base = pl.cast(local_base, pl.INT32) + local_token_count = pl.tensor.dim(local_request_ids, 0) + + with pl.spmd(1, name_hint="prefill_lower_local_request_ids"): + block_idx = pl.tile.get_block_idx() + if block_idx == 0: + for local_token in pl.range(local_token_count): + packed_token = base + pl.cast(local_token, pl.INT32) + pl.write(local_request_ids, [local_token], pl.cast(-1, pl.INT32)) + for request in pl.range(request_count): + request_start = pl.read(query_start_loc, [request]) + request_end = pl.read(query_start_loc, [request + 1]) + if packed_token >= request_start: + if packed_token < request_end: + pl.write(local_request_ids, [local_token], pl.cast(request, pl.INT32)) + return local_request_ids + + +@pl.jit +def prefill_metadata_test( + query_start_loc: pl.Tensor[[QUERY_START_LOC_DYN], pl.INT32], + local_base: pl.Scalar[pl.INT32], + request_ids: pl.Out[pl.Tensor[[LOCAL_TOKENS_DYN], pl.INT32]], +): + """Test rank-local packed prefill metadata lowering.""" + query_start_loc.bind_dynamic(0, QUERY_START_LOC_DYN) + request_ids.bind_dynamic(0, LOCAL_TOKENS_DYN) + return lower_local_request_ids(query_start_loc, request_ids, local_base) + + +def build_tensor_specs(): + import torch + from golden import ScalarSpec, TensorSpec + + query_start_loc = torch.tensor([0, 3, 7], dtype=torch.int32) + local_token_count = 4 + request_ids = torch.full((local_token_count,), -1, dtype=torch.int32) + local_base = local_token_count + return [ + TensorSpec("query_start_loc", list(query_start_loc.shape), torch.int32, init_value=query_start_loc), + ScalarSpec("local_base", torch.int32, local_base), + TensorSpec("request_ids", list(request_ids.shape), torch.int32, is_output=True), + ] + + +def golden_prefill_metadata(tensors): + import torch + + query_start_loc = tensors["query_start_loc"] + local_token_count = tensors["request_ids"].shape[0] + local_base = int(tensors["local_base"]) + total_tokens = int(query_start_loc[-1]) + request_ids = tensors["request_ids"] + request_ids.fill_(-1) + for local_token in range(local_token_count): + packed_token = local_base + local_token + if packed_token >= total_tokens: + continue + for request in range(query_start_loc.numel() - 1): + if query_start_loc[request] <= packed_token < query_start_loc[request + 1]: + request_ids[local_token] = request + break + expected = torch.tensor([1, 1, 1, -1], dtype=torch.int32) + if not torch.equal(request_ids, expected): + raise AssertionError(f"unexpected request ids: {request_ids.tolist()}") + + +if __name__ == "__main__": + import argparse + + from golden import run_jit + + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("-p", "--platform", default="a2a3", choices=["a2a3", "a2a3sim", "a5", "a5sim"]) + parser.add_argument("-d", "--device", type=int, default=0) + parser.add_argument("--compile-only", action="store_true") + args = parser.parse_args() + + result = run_jit( + fn=prefill_metadata_test, + specs=build_tensor_specs(), + golden_fn=golden_prefill_metadata, + compile_only=args.compile_only, + runtime_cfg={"platform": args.platform, "device_id": args.device}, + ) + if not result.passed: + if result.error: + print(result.error) + raise SystemExit(1) diff --git a/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py b/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py index 13dfcab3d..dbf89f8a6 100644 --- a/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py +++ b/models/deepseek_v4_flash_dspark/prefill_sparse_attn.py @@ -21,6 +21,7 @@ PREFILL_BATCH, PREFILL_SEQ, ) +from prefill_metadata import REQUESTS_DYN # Longest sequence the model config admits; the host-side bound on token_count. MAX_SEQ_LEN = M.max_position_embeddings @@ -217,6 +218,7 @@ def _hca_streaming_wave( rope_cos_il: pl.Tensor[[T_PAD, ROPE_DIM], pl.FP32], rope_sin_signed: pl.Tensor[[T_PAD, ROPE_DIM], pl.FP32], wave_completion: pl.Array[1, pl.TASK_ID], + request_offset: pl.Scalar[pl.INDEX], tile_base: pl.Scalar[pl.INDEX], dense_query_base: pl.Scalar[pl.INDEX], ): @@ -225,7 +227,8 @@ def _hca_streaming_wave( ori_block_num = pl.tensor.dim(ori_kv, 0) ori_cache_rows = ori_block_num * BLOCK_SIZE ori_kv_flat = pl.reshape(ori_kv, [ori_cache_rows, HEAD_DIM]) - query_base = tile_base + dense_query_base + query_base = request_offset + tile_base + dense_query_base + request_end = request_offset + active_rows with pl.spmd( HCA_QUERY_TILE // HCA_GATHER_TOKEN_TILE, @@ -239,7 +242,7 @@ def _hca_streaming_wave( gather_dst = gather_local_t * WIN raw_stage = pl.full([WIN, HEAD_DIM], dtype=pl.BF16, value=0.0) valid_stage = pl.full([1, WIN], dtype=pl.FP32, value=0.0) - if gather_t < active_rows: + if gather_t < request_end: for gather_k in pl.range(WIN): gather_row_i32 = pl.read(swa_indices, [gather_t, gather_k]) if gather_row_i32 >= 0: @@ -256,7 +259,7 @@ def _hca_streaming_wave( with pl.spmd(HCA_QUERY_TILE, name_hint="prefill_hca_stream_raw_qk_pv", deps=[raw_gather_tid]) as raw_heads_tid: raw_local_t = pl.tile.get_block_idx() raw_t = query_base + raw_local_t - if raw_t < active_rows: + if raw_t < request_end: raw_src = raw_local_t * WIN raw_kv_tile = raw_kv[raw_src : raw_src + WIN, 0:HEAD_DIM] raw_valid_row = raw_valid[raw_local_t : raw_local_t + 1, 0:WIN] @@ -295,7 +298,7 @@ def _hca_streaming_wave( qk_hb = qk_item - qk_local_t * (H // QK_M_TILE) qk_t = query_base + qk_local_t qk_token_base = qk_local_t * (H // HEAD_TILE) * HCA_CMP_WORK_COUNT * HEAD_TILE - if qk_t < active_rows: + if qk_t < request_end: qk_position_i32 = pl.read(position_ids, [qk_t]) if qk_position_i32 >= 0: qk_visible_rows = (qk_position_i32 + 1) // HCA_COMPRESS_RATIO @@ -350,7 +353,7 @@ def _hca_streaming_wave( merge_local_t = merge_item // (H // HEAD_TILE) merge_h_idx = merge_item - merge_local_t * (H // HEAD_TILE) merge_t = query_base + merge_local_t - if merge_t < active_rows: + if merge_t < request_end: merge_h0 = merge_h_idx * HEAD_TILE merge_stream_row = merge_local_t * H + merge_h0 merge_m = stream_state_m[merge_stream_row : merge_stream_row + HEAD_TILE, 0:1] @@ -425,6 +428,7 @@ def _hca_streaming_attn_tile( cmp_block_table: pl.Tensor[[HCA_CMP_MAX_BLOCKS], pl.INT32], position_ids: pl.Tensor[[T_DYN], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], + request_offset: pl.Scalar[pl.INDEX], active_rows: pl.Scalar[pl.INDEX], freqs_cos: pl.Tensor[[T_DYN, ROPE_DIM], pl.BF16], freqs_sin: pl.Tensor[[T_DYN, ROPE_DIM], pl.BF16], @@ -450,7 +454,7 @@ def _hca_streaming_attn_tile( level=pl.Level.CORE_GROUP, name_hint="prefill_hca_stream_cmp_plan", deps=[packed_init_tid], ) as cmp_plan_tid: # Maximum compressed-history rows for this tile. - plan_last_t = tile_base + tile_rows - 1 + plan_last_t = request_offset + tile_base + tile_rows - 1 plan_max_pos = pl.read(position_ids, [plan_last_t]) plan_visible_rows = (plan_max_pos + 1) // HCA_COMPRESS_RATIO if plan_visible_rows < 0: @@ -494,7 +498,7 @@ def _hca_streaming_attn_tile( rope_cs_tid = _prepare_sparse_attn_rope( freqs_cos, freqs_sin, rope_cos_il, rope_sin_signed, rope_swap_idx, - tile_base, tile_rows, + request_offset + tile_base, tile_rows, ) # Serial query-wave completion. @@ -518,7 +522,7 @@ def _hca_streaming_attn_tile( stream_state_m, stream_state_l, stream_heads, cmp_partial_m, cmp_partial_l, cmp_partial_o, rope_cos_il, rope_sin_signed, - wave_completion, tile_base, dense_query_base, + wave_completion, request_offset, tile_base, dense_query_base, ) heads_complete_tid = pl.system.task_dummy(deps=[wave_completion[0]]) @@ -528,7 +532,7 @@ def _hca_streaming_attn_tile( o_packed_heads, wo_a, wo_b, wo_b_scale, attn_out, - tile_base, tile_rows, + request_offset + tile_base, tile_rows, heads_complete_tid, o_proj_weight_dep, ) tile_completion[0] = act_tid @@ -540,7 +544,8 @@ def _sparse_attn_wave( ori_kv: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], swa_indices: pl.Tensor[[T_DYN, WIN], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, CMP_MAX_BLOCKS], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_indices: pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32], valid_block_mask: pl.Tensor[[T_DYN, VALID_BLOCK_MASK_COLS], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], @@ -584,13 +589,15 @@ def _sparse_attn_wave( gather_t = query_base + gather_local_t if gather_t < t_dim: if gather_t < active_rows: + request_id = pl.read(local_request_ids, [gather_t]) block_base = gather_local_t * PREFILL_SPARSE_PAD stage = pl.full([PREFILL_ATTN_TILE, HEAD_DIM], dtype=pl.BF16, value=0.0) - for gather_ki in pl.range(PREFILL_ATTN_TILE): - gather_raw = pl.read(swa_indices, [gather_t, gather_ki]) - if gather_raw >= 0: - src = pl.cast(gather_raw, pl.INDEX) - stage[gather_ki : gather_ki + 1, :] = ori_kv_flat[src : src + 1, :] + if request_id >= 0: + for gather_ki in pl.range(PREFILL_ATTN_TILE): + gather_raw = pl.read(swa_indices, [gather_t, gather_ki]) + if gather_raw >= 0: + src = pl.cast(gather_raw, pl.INDEX) + stage[gather_ki : gather_ki + 1, :] = ori_kv_flat[src : src + 1, :] sparse_kv[block_base : block_base + PREFILL_ATTN_TILE, :] = stage with pl.spmd( @@ -609,10 +616,11 @@ def _sparse_attn_wave( gather_t = query_base + gather_local_t if gather_t < t_dim: if gather_t < active_rows: + request_id = pl.read(local_request_ids, [gather_t]) gather_block_valid = pl.read(valid_block_mask, [gather_t, gather_sb]) - if gather_block_valid > 0: - block_base = gather_local_t * PREFILL_SPARSE_PAD + gather_k0 - stage = pl.full([PREFILL_ATTN_TILE, HEAD_DIM], dtype=pl.BF16, value=0.0) + block_base = gather_local_t * PREFILL_SPARSE_PAD + gather_k0 + stage = pl.full([PREFILL_ATTN_TILE, HEAD_DIM], dtype=pl.BF16, value=0.0) + if request_id >= 0 and gather_block_valid > 0: for gather_ki in pl.range(PREFILL_ATTN_TILE): gather_cmp_k = gather_k0 + gather_ki - WIN if gather_cmp_k < IDX_TOPK: @@ -620,10 +628,12 @@ def _sparse_attn_wave( if gather_raw >= 0: cmp_slot = gather_raw blk_slot = cmp_slot // BLOCK_SIZE - blk = pl.cast(pl.read(cmp_block_table, [blk_slot]), pl.INDEX) - src = blk * BLOCK_SIZE + (cmp_slot - blk_slot * BLOCK_SIZE) - stage[gather_ki : gather_ki + 1, :] = cmp_kv_flat[src : src + 1, :] - sparse_kv[block_base : block_base + PREFILL_ATTN_TILE, :] = stage + blk_raw = pl.read(cmp_block_table, [request_id, blk_slot]) + if blk_raw >= 0 and blk_raw < cmp_block_num: + blk = pl.cast(blk_raw, pl.INDEX) + src = blk * BLOCK_SIZE + (cmp_slot - blk_slot * BLOCK_SIZE) + stage[gather_ki : gather_ki + 1, :] = cmp_kv_flat[src : src + 1, :] + sparse_kv[block_base : block_base + PREFILL_ATTN_TILE, :] = stage # Keep the existing 16-row vectorized bias path inside each query wave. with pl.spmd( @@ -832,7 +842,8 @@ def _sparse_attn_heads( ori_kv: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], swa_indices: pl.Tensor[[T_DYN, WIN], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, CMP_MAX_BLOCKS], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_indices: pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32], valid_block_mask: pl.Tensor[[T_DYN, VALID_BLOCK_MASK_COLS], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], @@ -885,6 +896,7 @@ def _sparse_attn_heads( swa_indices, cmp_kv, cmp_block_table, + local_request_ids, cmp_indices, valid_block_mask, attn_sink, @@ -1091,9 +1103,10 @@ def hca_streaming_attn_physical( attn_out: pl.Tensor[[T_DYN, D], pl.BF16], cache_ready_dep: pl.Scalar[pl.TASK_ID], o_proj_weight_dep: pl.Scalar[pl.TASK_ID], + request_offset: pl.Scalar[pl.INDEX], + active_rows: pl.Scalar[pl.INDEX], ): """Run ratio-128 attention over streamed history.""" - active_rows = pl.tensor.dim(q, 0) tile_completion = pl.array.create(1, pl.TASK_ID) tile_completion[0] = cache_ready_dep for tile_base in pl.range(0, active_rows, T_PAD): @@ -1113,7 +1126,7 @@ def hca_streaming_attn_physical( _hca_streaming_attn_tile( q, ori_kv, swa_indices, cmp_kv, cmp_block_table, - position_ids, attn_sink, active_rows, + position_ids, attn_sink, request_offset, active_rows, freqs_cos, freqs_sin, wo_a, wo_b, wo_b_scale, attn_out, @@ -1121,10 +1134,12 @@ def hca_streaming_attn_physical( tile_completion, tile_base, tile_rows, ) # Publish final o-proj completion outside the manual tile scope. - with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_hca_stream_publish", deps=[tile_completion[0]]): - completion_anchor = pl.read(attn_out, [0, 0]) - pl.write(attn_out, [0, 0], completion_anchor) - return attn_out + with pl.at( + level=pl.Level.CORE_GROUP, name_hint="prefill_hca_stream_publish", deps=[tile_completion[0]] + ) as publish_tid: + completion_anchor = pl.read(attn_out, [request_offset, 0]) + pl.write(attn_out, [request_offset, 0], completion_anchor) + return publish_tid @pl.jit.inline(auto_scope=False) @@ -1133,7 +1148,8 @@ def sparse_attn_compute( ori_kv: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], swa_indices: pl.Tensor[[T_DYN, WIN], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, CMP_MAX_BLOCKS], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_indices: pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32], valid_block_mask: pl.Tensor[[T_DYN, VALID_BLOCK_MASK_COLS], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], @@ -1166,6 +1182,7 @@ def sparse_attn_compute( swa_indices, cmp_kv, cmp_block_table, + local_request_ids, cmp_indices, valid_block_mask, attn_sink, @@ -1197,7 +1214,8 @@ def sparse_attn_physical( ori_kv: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], swa_indices: pl.Tensor[[T_DYN, WIN], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, CMP_MAX_BLOCKS], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_indices: pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32], valid_block_mask: pl.Tensor[[T_DYN, VALID_BLOCK_MASK_COLS], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], @@ -1217,6 +1235,7 @@ def sparse_attn_physical( swa_indices, cmp_kv, cmp_block_table, + local_request_ids, cmp_indices, valid_block_mask, attn_sink, @@ -1237,7 +1256,8 @@ def prefill_sparse_attn_test( ori_kv: pl.Tensor[[ORI_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], swa_indices: pl.Tensor[[T_DYN, WIN], pl.INT32], cmp_kv: pl.Tensor[[CMP_BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - cmp_block_table: pl.Tensor[[CMP_MAX_BLOCKS], pl.INT32], + cmp_block_table: pl.Tensor[[REQUESTS_DYN, CMP_MAX_BLOCKS], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], cmp_indices: pl.Tensor[[T_DYN, IDX_TOPK], pl.INT32], valid_block_mask: pl.Tensor[[T_DYN, VALID_BLOCK_MASK_COLS], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], @@ -1250,9 +1270,11 @@ def prefill_sparse_attn_test( ): ori_kv.bind_dynamic(0, ORI_BLOCK_NUM_DYN) cmp_kv.bind_dynamic(0, CMP_BLOCK_NUM_DYN) + cmp_block_table.bind_dynamic(0, REQUESTS_DYN) q.bind_dynamic(0, T_DYN) swa_indices.bind_dynamic(0, T_DYN) cmp_indices.bind_dynamic(0, T_DYN) + local_request_ids.bind_dynamic(0, T_DYN) valid_block_mask.bind_dynamic(0, T_DYN) freqs_cos.bind_dynamic(0, T_DYN) freqs_sin.bind_dynamic(0, T_DYN) @@ -1264,6 +1286,7 @@ def prefill_sparse_attn_test( swa_indices, cmp_kv, cmp_block_table, + local_request_ids, cmp_indices, valid_block_mask, attn_sink, @@ -1286,6 +1309,7 @@ def golden_prefill_sparse_attn(tensors): ori_kv = tensors["ori_kv"].float() cmp_kv = tensors["cmp_kv"].float() cmp_block_table = tensors["cmp_block_table"] + local_request_ids = tensors["local_request_ids"] swa_indices = tensors["swa_indices"] cmp_indices = tensors["cmp_indices"] attn_sink = tensors["attn_sink"].float() @@ -1297,6 +1321,9 @@ def golden_prefill_sparse_attn(tensors): o = torch.zeros(token_count, H, HEAD_DIM) for t in range(token_count): + request_id = int(local_request_ids[t].item()) + if request_id < 0: + continue gathered = [] for row_i in swa_indices[t].tolist(): row = int(row_i) @@ -1307,7 +1334,7 @@ def golden_prefill_sparse_attn(tensors): cmp_slot = int(raw_i) if cmp_slot < 0 or cmp_slot >= CMP_MAX_BLOCKS * BLOCK_SIZE: continue - block_id = int(cmp_block_table[cmp_slot // BLOCK_SIZE].item()) + block_id = int(cmp_block_table[request_id, cmp_slot // BLOCK_SIZE].item()) intra = cmp_slot % BLOCK_SIZE if block_id >= 0: gathered.append(cmp_kv[block_id, intra, 0]) @@ -1414,11 +1441,14 @@ def init_cmp_kv(): return ((torch.rand(cmp_block_num, BLOCK_SIZE, 1, HEAD_DIM) - 0.5) * 0.05).to(torch.bfloat16) def init_cmp_block_table(): - table = torch.zeros(CMP_MAX_BLOCKS, dtype=torch.int32) + table = torch.zeros(1, CMP_MAX_BLOCKS, dtype=torch.int32) for blk in range(CMP_MAX_BLOCKS): - table[blk] = blk % cmp_block_num + table[0, blk] = blk % cmp_block_num return table + def init_local_request_ids(): + return torch.zeros(token_count, dtype=torch.int32) + def init_swa_indices(): idx = torch.full((token_count, WIN), -1, dtype=torch.int32) for t in range(token_count): @@ -1484,7 +1514,8 @@ def init_wo_b(): TensorSpec( "cmp_kv", [cmp_block_num, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_cmp_kv ), - TensorSpec("cmp_block_table", [CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), + TensorSpec("cmp_block_table", [1, CMP_MAX_BLOCKS], torch.int32, init_value=init_cmp_block_table), + TensorSpec("local_request_ids", [token_count], torch.int32, init_value=init_local_request_ids), TensorSpec("cmp_indices", [token_count, IDX_TOPK], torch.int32, init_value=init_cmp_indices), TensorSpec( "valid_block_mask", diff --git a/models/deepseek_v4_flash_dspark/prefill_swa.py b/models/deepseek_v4_flash_dspark/prefill_swa.py index 58e2c5a87..9a8246f19 100644 --- a/models/deepseek_v4_flash_dspark/prefill_swa.py +++ b/models/deepseek_v4_flash_dspark/prefill_swa.py @@ -31,6 +31,7 @@ materialize_spec, prefill_cp_token_allgather_step, ) +from prefill_metadata import REQUESTS_DYN from prefill_o_proj import ( O_PROJ_LOCAL_COLS, O_PROJ_LOCAL_GROUPS, @@ -112,9 +113,10 @@ def prefill_attention_swa( freqs_cos: pl.Tensor[[T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[T_DYN, ROPE_HEAD_DIM], pl.BF16], kv_cache: pl.Tensor[[BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - block_table: pl.Tensor[[BLOCK_TABLE_BLOCKS], pl.INT32], + block_table: pl.Tensor[[REQUESTS_DYN, BLOCK_TABLE_BLOCKS], pl.INT32], ori_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], position_ids: pl.Tensor[[T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], @@ -158,37 +160,39 @@ def prefill_attention_swa( for idx_t in pl.range(t_dim): idx_row = pl.full([1, WIN], dtype=pl.INT32, value=-1) mask_row = pl.full([1, VALID_BLOCK_MASK_COLS], dtype=pl.INT32, value=0) - abs_pos = pl.read(position_ids, [idx_t]) - window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) - key_start_abs = abs_pos + 1 - window_valid - for win_col in pl.range(WIN): - win_col_i32 = pl.cast(win_col, pl.INT32) - if win_col_i32 < window_valid: - key_abs = key_start_abs + win_col_i32 - blk_slot = key_abs // BLOCK_SIZE - blk = pl.read(block_table, [pl.cast(blk_slot, pl.INDEX)]) - if blk >= 0: - row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) - pl.write(idx_row, [0, win_col], row) - if win_col < SPARSE_BIAS_COLS: - block_col = win_col // PREFILL_ATTN_TILE - pl.write(mask_row, [0, block_col], pl.cast(1, pl.INT32)) + request_id = pl.read(local_request_ids, [idx_t]) + if request_id >= 0: + abs_pos = pl.read(position_ids, [idx_t]) + window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) + key_start_abs = abs_pos + 1 - window_valid + for win_col in pl.range(WIN): + win_col_i32 = pl.cast(win_col, pl.INT32) + if win_col_i32 < window_valid: + key_abs = key_start_abs + win_col_i32 + blk_slot = key_abs // BLOCK_SIZE + blk = pl.read(block_table, [request_id, pl.cast(blk_slot, pl.INDEX)]) + if blk >= 0: + row = pl.cast(blk * BLOCK_SIZE + (key_abs - blk_slot * BLOCK_SIZE), pl.INT32) + pl.write(idx_row, [0, win_col], row) + if win_col < SPARSE_BIAS_COLS: + block_col = win_col // PREFILL_ATTN_TILE + pl.write(mask_row, [0, block_col], pl.cast(1, pl.INT32)) swa_indices[idx_t : idx_t + 1, 0:WIN] = idx_row valid_block_mask[idx_t : idx_t + 1, 0:VALID_BLOCK_MASK_COLS] = mask_row - cmp_block_table_dummy = pl.create_tensor([SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32) + request_count = pl.tensor.dim(block_table, 0) + cmp_block_table_dummy = pl.create_tensor([request_count, SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32) cmp_kv_dummy = pl.create_tensor([CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], dtype=pl.BF16) cmp_indices_dummy = pl.create_tensor([t_dim, IDX_TOPK], dtype=pl.INT32) - cmp_block_table_dummy_2d = pl.reshape(cmp_block_table_dummy, [1, SPARSE_CMP_MAX_BLOCKS]) - with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_swa_cmp_dummy_init"): - cmp_block_table_dummy_2d[:, :] = pl.full([1, SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32, value=0) + for request in pl.spmd(request_count, name_hint="prefill_swa_cmp_dummy_init"): + cmp_block_table_dummy[request : request + 1, :] = pl.full([1, SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32, value=0) for dummy_t in pl.spmd(t_dim, name_hint="prefill_swa_cmp_indices_dummy_init"): cmp_indices_dummy[dummy_t : dummy_t + 1, :] = pl.full([1, IDX_TOPK], dtype=pl.INT32, value=-1) attn_out = pl.create_tensor([t_dim, D], dtype=pl.BF16) o_proj_weight_dep = pl.system.task_dummy(deps=[]) attn_out = sparse_attn_physical( q, kv_cache, swa_indices, - cmp_kv_dummy, cmp_block_table_dummy, + cmp_kv_dummy, cmp_block_table_dummy, local_request_ids, cmp_indices_dummy, valid_block_mask, attn_sink, @@ -216,9 +220,10 @@ def prefill_attention_swa_test( freqs_cos: pl.Tensor[[T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[T_DYN, ROPE_HEAD_DIM], pl.BF16], kv_cache: pl.InOut[pl.Tensor[[BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - block_table: pl.Tensor[[BLOCK_TABLE_BLOCKS], pl.INT32], + block_table: pl.Tensor[[REQUESTS_DYN, BLOCK_TABLE_BLOCKS], pl.INT32], ori_slot_mapping: pl.Tensor[[T_DYN], pl.INT64], position_ids: pl.Tensor[[T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[T_DYN], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], @@ -229,8 +234,10 @@ def prefill_attention_swa_test( freqs_cos.bind_dynamic(0, T_DYN) freqs_sin.bind_dynamic(0, T_DYN) kv_cache.bind_dynamic(0, BLOCK_NUM_DYN) + block_table.bind_dynamic(0, REQUESTS_DYN) ori_slot_mapping.bind_dynamic(0, T_DYN) position_ids.bind_dynamic(0, T_DYN) + local_request_ids.bind_dynamic(0, T_DYN) x_out.bind_dynamic(0, T_DYN) prefill_attention_swa( @@ -239,7 +246,7 @@ def prefill_attention_swa_test( attn_norm_w, wq_a, wq_b, wq_b_scale, wkv, gamma_cq, gamma_ckv, freqs_cos, freqs_sin, kv_cache, block_table, ori_slot_mapping, - position_ids, + position_ids, local_request_ids, attn_sink, wo_a, wo_b, wo_b_scale, x_out, ) @@ -261,10 +268,11 @@ def prefill_attention_swa_cp_core( freqs_cos_full: pl.Tensor[[CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin_full: pl.Tensor[[CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], kv_cache: pl.Tensor[[BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - block_table: pl.Tensor[[BLOCK_TABLE_BLOCKS], pl.INT32], + block_table: pl.Tensor[[REQUESTS_DYN, BLOCK_TABLE_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], wo_a: pl.Tensor[[O_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_GROUPS * O_LORA], pl.INT8], @@ -310,36 +318,38 @@ def prefill_attention_swa_cp_core( for idx_t in pl.range(q_dim): idx_row = pl.full([1, WIN], dtype=pl.INT32, value=-1) mask_row = pl.full([1, VALID_BLOCK_MASK_COLS], dtype=pl.INT32, value=0) - abs_pos = pl.read(position_ids_local, [idx_t]) - window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) - key_start_abs = abs_pos + 1 - window_valid - for win_col in pl.range(WIN): - win_col_i32 = pl.cast(win_col, pl.INT32) - if win_col_i32 < window_valid: - key_abs = key_start_abs + win_col_i32 - blk_slot = key_abs // BLOCK_SIZE - blk = pl.read(block_table, [pl.cast(blk_slot, pl.INDEX)]) - if blk >= 0: - block_row = key_abs - blk_slot * BLOCK_SIZE - row = pl.cast(blk * BLOCK_SIZE + block_row, pl.INT32) - pl.write(idx_row, [0, win_col], row) - if win_col < SPARSE_BIAS_COLS: - block_col = win_col // PREFILL_ATTN_TILE - pl.write(mask_row, [0, block_col], pl.cast(1, pl.INT32)) + request_id = pl.read(local_request_ids, [idx_t]) + if request_id >= 0: + abs_pos = pl.read(position_ids_local, [idx_t]) + window_valid = pl.min(pl.cast(WIN, pl.INT32), abs_pos + 1) + key_start_abs = abs_pos + 1 - window_valid + for win_col in pl.range(WIN): + win_col_i32 = pl.cast(win_col, pl.INT32) + if win_col_i32 < window_valid: + key_abs = key_start_abs + win_col_i32 + blk_slot = key_abs // BLOCK_SIZE + blk = pl.read(block_table, [request_id, pl.cast(blk_slot, pl.INDEX)]) + if blk >= 0: + block_row = key_abs - blk_slot * BLOCK_SIZE + row = pl.cast(blk * BLOCK_SIZE + block_row, pl.INT32) + pl.write(idx_row, [0, win_col], row) + if win_col < SPARSE_BIAS_COLS: + block_col = win_col // PREFILL_ATTN_TILE + pl.write(mask_row, [0, block_col], pl.cast(1, pl.INT32)) swa_indices[idx_t : idx_t + 1, 0:WIN] = idx_row valid_block_mask[idx_t : idx_t + 1, 0:VALID_BLOCK_MASK_COLS] = mask_row - cmp_block_table_dummy = pl.create_tensor([SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32) + request_count = pl.tensor.dim(block_table, 0) + cmp_block_table_dummy = pl.create_tensor([request_count, SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32) cmp_kv_dummy = pl.create_tensor([CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], dtype=pl.BF16) cmp_indices_dummy = pl.create_tensor([q_dim, IDX_TOPK], dtype=pl.INT32) - cmp_block_table_dummy_2d = pl.reshape(cmp_block_table_dummy, [1, SPARSE_CMP_MAX_BLOCKS]) - with pl.at(level=pl.Level.CORE_GROUP, name_hint="prefill_swa_cp_cmp_dummy_init"): - cmp_block_table_dummy_2d[:, :] = pl.full([1, SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32, value=0) + for request in pl.spmd(request_count, name_hint="prefill_swa_cp_cmp_dummy_init"): + cmp_block_table_dummy[request : request + 1, :] = pl.full([1, SPARSE_CMP_MAX_BLOCKS], dtype=pl.INT32, value=0) for dummy_t in pl.spmd(q_dim, name_hint="prefill_swa_cp_cmp_indices_dummy_init"): cmp_indices_dummy[dummy_t : dummy_t + 1, :] = pl.full([1, IDX_TOPK], dtype=pl.INT32, value=-1) attn_out_local = sparse_attn_physical( q, kv_cache, swa_indices, - cmp_kv_dummy, cmp_block_table_dummy, cmp_indices_dummy, + cmp_kv_dummy, cmp_block_table_dummy, local_request_ids, cmp_indices_dummy, valid_block_mask, attn_sink, freqs_cos_local, freqs_sin_local, wo_a, wo_b, wo_b_scale, @@ -364,10 +374,11 @@ def prefill_attention_swa_cp( freqs_cos: pl.Tensor[[CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], kv_cache: pl.Tensor[[BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16], - block_table: pl.Tensor[[BLOCK_TABLE_BLOCKS], pl.INT32], + block_table: pl.Tensor[[REQUESTS_DYN, BLOCK_TABLE_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], wo_a_local: pl.Tensor[[O_PROJ_LOCAL_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b_local: pl.Tensor[[D, O_PROJ_LOCAL_COLS], pl.INT8], @@ -428,7 +439,7 @@ def prefill_attention_swa_cp( freqs_cos, freqs_sin, kv_cache, block_table, ori_slot_mapping_full, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, attn_sink, wo_a_full, wo_b_full, wo_b_scale, attn_out_local, @@ -462,10 +473,11 @@ def prefill_attention_swa_cp_test( freqs_cos: pl.Tensor[[CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], kv_cache: pl.InOut[pl.Tensor[[BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - block_table: pl.Tensor[[BLOCK_TABLE_BLOCKS], pl.INT32], + block_table: pl.Tensor[[REQUESTS_DYN, BLOCK_TABLE_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[CP_KV_T_DYN], pl.INT64], position_ids_local: pl.Tensor[[CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[CP_Q_T_DYN], pl.INT32], attn_sink: pl.Tensor[[H], pl.FP32], wo_a: pl.Tensor[[O_PROJ_LOCAL_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[D, O_PROJ_LOCAL_COLS], pl.INT8], @@ -485,9 +497,11 @@ def prefill_attention_swa_cp_test( freqs_cos.bind_dynamic(0, CP_KV_T_DYN) freqs_sin.bind_dynamic(0, CP_KV_T_DYN) kv_cache.bind_dynamic(0, BLOCK_NUM_DYN) + block_table.bind_dynamic(0, REQUESTS_DYN) ori_slot_mapping_full.bind_dynamic(0, CP_KV_T_DYN) position_ids_local.bind_dynamic(0, CP_Q_T_DYN) position_ids_full.bind_dynamic(0, CP_KV_T_DYN) + local_request_ids.bind_dynamic(0, CP_Q_T_DYN) x_out_full.bind_dynamic(0, CP_KV_T_DYN) wo_a_full = pl.create_tensor([O_PROJ_SCRATCH_GROUPS, O_PROJ_SCRATCH_RANK, O_PROJ_SCRATCH_INPUT], dtype=pl.BF16) @@ -501,7 +515,7 @@ def prefill_attention_swa_cp_test( attn_norm_w, wq_a, wq_b, wq_b_scale, wkv, gamma_cq, gamma_ckv, freqs_cos, freqs_sin, kv_cache, block_table, ori_slot_mapping_full, - position_ids_local, position_ids_full, + position_ids_local, position_ids_full, local_request_ids, attn_sink, wo_a, wo_b, wo_b_scale, wo_a_full, wo_b_full, x_out_full, @@ -530,10 +544,11 @@ def l3_prefill_attention_swa_cp( freqs_cos: pl.Tensor[[TP_SIZE, CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], freqs_sin: pl.Tensor[[TP_SIZE, CP_KV_T_DYN, ROPE_HEAD_DIM], pl.BF16], kv_cache: pl.InOut[pl.Tensor[[TP_SIZE, BLOCK_NUM_DYN, BLOCK_SIZE, 1, HEAD_DIM], pl.BF16]], - block_table: pl.Tensor[[TP_SIZE, BLOCK_TABLE_BLOCKS], pl.INT32], + block_table: pl.Tensor[[TP_SIZE, REQUESTS_DYN, BLOCK_TABLE_BLOCKS], pl.INT32], ori_slot_mapping_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT64], position_ids_local: pl.Tensor[[TP_SIZE, CP_Q_T_DYN], pl.INT32], position_ids_full: pl.Tensor[[TP_SIZE, CP_KV_T_DYN], pl.INT32], + local_request_ids: pl.Tensor[[TP_SIZE, CP_Q_T_DYN], pl.INT32], attn_sink: pl.Tensor[[TP_SIZE, H], pl.FP32], wo_a: pl.Tensor[[TP_SIZE, O_PROJ_LOCAL_GROUPS, O_LORA, O_GROUP_IN], pl.BF16], wo_b: pl.Tensor[[TP_SIZE, D, O_PROJ_LOCAL_COLS], pl.INT8], @@ -545,9 +560,11 @@ def l3_prefill_attention_swa_cp( freqs_cos.bind_dynamic(1, CP_KV_T_DYN) freqs_sin.bind_dynamic(1, CP_KV_T_DYN) kv_cache.bind_dynamic(1, BLOCK_NUM_DYN) + block_table.bind_dynamic(1, REQUESTS_DYN) ori_slot_mapping_full.bind_dynamic(1, CP_KV_T_DYN) position_ids_local.bind_dynamic(1, CP_Q_T_DYN) position_ids_full.bind_dynamic(1, CP_KV_T_DYN) + local_request_ids.bind_dynamic(1, CP_Q_T_DYN) x_out_full.bind_dynamic(1, CP_KV_T_DYN) gather_window_buf = pld.alloc_window_buffer([PREFILL_GROUP_CAP, D], dtype=pl.BF16) @@ -575,7 +592,7 @@ def l3_prefill_attention_swa_cp( wkv[rank], gamma_cq[rank], gamma_ckv[rank], freqs_cos[rank], freqs_sin[rank], kv_cache[rank], block_table[rank], ori_slot_mapping_full[rank], - position_ids_local[rank], position_ids_full[rank], + position_ids_local[rank], position_ids_full[rank], local_request_ids[rank], attn_sink[rank], wo_a[rank], wo_b[rank], wo_b_scale[rank], x_out_full[rank], gather_window, gather_signal, @@ -654,12 +671,16 @@ def build_swa_metadata(): idx = torch.full((token_count, WIN), -1, dtype=torch.int32) pos = tensors["position_ids"] table = tensors["block_table"] + request_ids = tensors["local_request_ids"] for t in range(token_count): + request_id = int(request_ids[t].item()) + if request_id < 0: + continue abs_pos = int(pos[t].item()) window_valid = min(WIN, abs_pos + 1) key_start_abs = abs_pos + 1 - window_valid for k, key_abs in enumerate(range(key_start_abs, abs_pos + 1)): - row = cache_row_from_table(table, key_abs) + row = cache_row_from_table(table[request_id], key_abs) if row >= 0: idx[t, k] = row return idx @@ -670,7 +691,10 @@ def build_swa_metadata(): "ori_kv": kv_cache_in, "swa_indices": build_swa_metadata(), "cmp_kv": torch.zeros(CMP_BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.bfloat16), - "cmp_block_table": torch.zeros(SPARSE_CMP_MAX_BLOCKS, dtype=torch.int32), + "cmp_block_table": torch.zeros( + tensors["block_table"].shape[0], SPARSE_CMP_MAX_BLOCKS, dtype=torch.int32 + ), + "local_request_ids": tensors["local_request_ids"], "cmp_indices": torch.full((token_count, IDX_TOPK), -1, dtype=torch.int32), "attn_sink": tensors["attn_sink"], "freqs_cos": rope_cos_t, @@ -757,9 +781,9 @@ def init_freqs_cos(): def init_freqs_sin(): return shared_freqs_sin.clone() def init_block_table(): - tbl = torch.full((BLOCK_TABLE_BLOCKS,), -1, dtype=torch.int32) + tbl = torch.full((1, BLOCK_TABLE_BLOCKS), -1, dtype=torch.int32) for block in range(BLOCK_TABLE_BLOCKS): - tbl[block] = block % BLOCK_NUM + tbl[0, block] = block % BLOCK_NUM return tbl def init_kv_cache(): cache = torch.zeros(BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM) @@ -767,7 +791,7 @@ def init_kv_cache(): table = init_block_table() start = max(0, context_len - WIN) for abs_pos in range(start, context_len): - row = cache_row_from_table(table, abs_pos) + row = cache_row_from_table(table[0], abs_pos) value = (torch.rand(HEAD_DIM,) - 0.5) * 0.1 if row >= 0: cache_flat[row] = value.to(torch.bfloat16) @@ -777,10 +801,12 @@ def init_ori_slot_mapping(): pos = token_pos() table = init_block_table() for t in range(token_count): - mapping[t] = cache_row_from_table(table, int(pos[t].item())) + mapping[t] = cache_row_from_table(table[0], int(pos[t].item())) return mapping def init_position_ids(): return token_pos() + def init_local_request_ids(): + return torch.zeros(token_count, dtype=torch.int32) def init_attn_sink(): return torch.zeros(H) def init_wo_a(): @@ -809,9 +835,10 @@ def init_wo_b(): TensorSpec("freqs_sin", [token_count, ROPE_HEAD_DIM], torch.bfloat16, init_value=init_freqs_sin), TensorSpec("kv_cache", [BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM], torch.bfloat16, init_value=init_kv_cache, is_output=True), - TensorSpec("block_table", [BLOCK_TABLE_BLOCKS], torch.int32, init_value=init_block_table), + TensorSpec("block_table", [1, BLOCK_TABLE_BLOCKS], torch.int32, init_value=init_block_table), TensorSpec("ori_slot_mapping", [token_count], torch.int64, init_value=init_ori_slot_mapping), TensorSpec("position_ids", [token_count], torch.int32, init_value=init_position_ids), + TensorSpec("local_request_ids", [token_count], torch.int32, init_value=init_local_request_ids), TensorSpec("attn_sink", [H], torch.float32, init_value=init_attn_sink), TensorSpec("wo_a", [O_GROUPS, O_LORA, O_GROUP_IN], torch.bfloat16, init_value=init_wo_a), TensorSpec("wo_b", [D, O_GROUPS * O_LORA], torch.int8, init_value=lambda: wo_b_i8), @@ -859,6 +886,11 @@ def build_cp_tensor_specs( specs.append(TensorSpec( "position_ids_full", [tp_size, token_count], spec.dtype, init_value=cp_stack(value, tp_size), )) + elif spec.name == "local_request_ids": + specs.append(TensorSpec( + "local_request_ids", [tp_size, local_t], spec.dtype, + init_value=torch.zeros(tp_size, local_t, dtype=spec.dtype), + )) elif spec.name == "wo_a": shards = [value[rank * O_PROJ_LOCAL_GROUPS : (rank + 1) * O_PROJ_LOCAL_GROUPS] for rank in range(tp_size)] specs.append(TensorSpec( @@ -881,6 +913,70 @@ def build_cp_tensor_specs( return specs +def build_ragged2_cp_tensor_specs(tp_size: int = TP_SIZE): + """Build the two-request rank-crossing CP fixture from the B1 specs.""" + import torch + + from golden import TensorSpec + from utils import ( + block_table as make_block_table, + cache_row_from_table, + ori_slot_mapping as make_ori_slot_mapping, + token_local_rope, + ) + + if tp_size != 2: + raise ValueError(f"ragged2 requires tp_size=2, got {tp_size}") + + token_count = 8 + request_positions = ( + torch.tensor([126, 127, 128], dtype=torch.int32), + torch.tensor([30, 31, 32, 33], dtype=torch.int32), + ) + position_ids = torch.cat((*request_positions, torch.zeros(1, dtype=torch.int32))) + request_ids = torch.tensor([0, 0, 0, 1, 1, 1, 1, -1], dtype=torch.int32) + table = make_block_table(batch=2, table_blocks=BLOCK_TABLE_BLOCKS, physical_blocks=BLOCK_NUM) + + active_slot_mapping = [] + for request, positions in enumerate(request_positions): + request_table = table[request : request + 1] + request_mapping = make_ori_slot_mapping(positions.unsqueeze(0), request_table) + active_slot_mapping.append(request_mapping.reshape(-1)) + ori_slot_mapping = torch.cat((*active_slot_mapping, torch.full((1,), -1, dtype=torch.int64))) + + kv_cache = torch.zeros(BLOCK_NUM, BLOCK_SIZE, 1, HEAD_DIM, dtype=torch.bfloat16) + kv_cache_flat = kv_cache.view(BLOCK_NUM * BLOCK_SIZE, HEAD_DIM) + for request, start_pos in enumerate((126, 30)): + for position in range(max(0, start_pos - WIN), start_pos): + row = cache_row_from_table(table[request], position) + kv_cache_flat[row] = ((torch.rand(HEAD_DIM) - 0.5) * 0.1).to(torch.bfloat16) + + freqs_cos, freqs_sin = token_local_rope(M, 0, position_ids, max_seq_len=MAX_SEQ_LEN, dtype=torch.bfloat16) + replacements = { + "freqs_cos": cp_stack(freqs_cos, tp_size), + "freqs_sin": cp_stack(freqs_sin, tp_size), + "kv_cache": cp_stack(kv_cache, tp_size), + "block_table": cp_stack(table, tp_size), + "ori_slot_mapping_full": cp_stack(ori_slot_mapping, tp_size), + "position_ids_local": position_ids.reshape(tp_size, token_count // tp_size).contiguous(), + "position_ids_full": cp_stack(position_ids, tp_size), + "local_request_ids": request_ids.reshape(tp_size, token_count // tp_size).contiguous(), + } + + specs = [] + for spec in build_cp_tensor_specs(start_pos=0, token_count=token_count, tp_size=tp_size): + value = replacements.get(spec.name) + if value is None: + specs.append(spec) + continue + replacement_spec = TensorSpec( + spec.name, list(value.shape), spec.dtype, init_value=value, + is_output=spec.is_output, resident=spec.resident, + ) + specs.append(replacement_spec) + return specs + + def golden_prefill_attention_swa_cp(tensors): """Single-die reference replicated across DSA-CP ranks.""" import torch @@ -904,6 +1000,9 @@ def golden_prefill_attention_swa_cp(tensors): "kv_cache": tensors["kv_cache"][0].clone(), "block_table": tensors["block_table"][0], "ori_slot_mapping": tensors["ori_slot_mapping_full"][0], + "local_request_ids": torch.cat( + [tensors["local_request_ids"][rank] for rank in range(tp_size)] + ), "position_ids": tensors["position_ids_full"][0], "attn_sink": tensors["attn_sink"][0], "wo_a": torch.cat([tensors["wo_a"][rank] for rank in range(tp_size)], dim=0), @@ -933,8 +1032,14 @@ def golden_prefill_attention_swa_cp(tensors): parser.add_argument("--compile-only", action="store_true", default=False) parser.add_argument("--start-pos", type=int, default=START_POS, help="Absolute position of the first physical query token.") - parser.add_argument("--token-count", "--num-tokens", dest="token_count", type=int, default=PREFILL_SEQ, - help="Physical query-token extent across the group; must divide by --tp.") + parser.add_argument( + "--token-count", "--num-tokens", dest="token_count", type=int, default=None, + help=f"B1 physical query-token extent across the group; defaults to {PREFILL_SEQ}. ragged2 is fixed at 8.", + ) + parser.add_argument( + "--case", choices=["b1", "ragged2"], default="b1", + help="Fixture case; ragged2 is the fixed two-request TP2 boundary case.", + ) parser.add_argument("--enable-chip-swimlane", action="store_true", default=False) parser.add_argument("--enable-dep-gen", action="store_true", default=False) parser.add_argument("--dump-passes", action="store_true", default=False) @@ -945,7 +1050,15 @@ def golden_prefill_attention_swa_cp(tensors): device_ids = [int(device) for device in args.device.split(",")] if len(device_ids) != TP_SIZE: parser.error(f"need exactly {TP_SIZE} devices, got {device_ids}") - if args.token_count % TP_SIZE != 0: + if args.case == "ragged2" and TP_SIZE != 2: + parser.error("--case ragged2 requires --tp 2") + if args.case == "ragged2" and args.start_pos != 0: + parser.error("--case ragged2 has fixed request starts and requires --start-pos 0") + if args.token_count is None: + args.token_count = 8 if args.case == "ragged2" else PREFILL_SEQ + if args.case == "ragged2" and args.token_count != 8: + parser.error("--case ragged2 has a fixed physical extent and requires --token-count 8") + if args.case == "b1" and args.token_count % TP_SIZE != 0: parser.error(f"--token-count must be a multiple of --tp={TP_SIZE}, got {args.token_count}") if TP_SIZE == 1: @@ -971,9 +1084,14 @@ def golden_prefill_attention_swa_cp(tensors): else: from pypto.ir.distributed_compiled_program import DistributedConfig + specs = ( + build_ragged2_cp_tensor_specs(TP_SIZE) + if args.case == "ragged2" + else build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE) + ) result = run_jit( fn=l3_prefill_attention_swa_cp, - specs=build_cp_tensor_specs(args.start_pos, args.token_count, TP_SIZE), + specs=specs, golden_fn=golden_prefill_attention_swa_cp, compile_cfg=dict( dump_passes=args.dump_passes,