Generate sparse block mask via top-k policy denoising. Fuses per-row budget (3-regime k_schedule + linear decay, both keyed on the full prompt KV length so the result is chunked-prefill invariant), radix top-k threshold, and fixed retention (initial / window / diagonal). Args:
(
block_logits: Tensor,
q_seq_lens: Tensor,
kv_seq_lens: Tensor,
num_prompt_tokens: Tensor,
block_size: int = 128,
alpha: float = 1.0,
initial_blocks: int = 4,
window_size: int = 4,
k_block_num_rate_medium: float = 0.2,
k_block_num_bias_medium: int = 30,
k_block_num_rate_large: float = 0.1,
k_block_num_bias_large: int = 30,
)
| 171 | |
| 172 | |
| 173 | def stem_tpd( |
| 174 | block_logits: Tensor, |
| 175 | q_seq_lens: Tensor, |
| 176 | kv_seq_lens: Tensor, |
| 177 | num_prompt_tokens: Tensor, |
| 178 | block_size: int = 128, |
| 179 | alpha: float = 1.0, |
| 180 | initial_blocks: int = 4, |
| 181 | window_size: int = 4, |
| 182 | k_block_num_rate_medium: float = 0.2, |
| 183 | k_block_num_bias_medium: int = 30, |
| 184 | k_block_num_rate_large: float = 0.1, |
| 185 | k_block_num_bias_large: int = 30, |
| 186 | ) -> Tensor: |
| 187 | """Generate sparse block mask via top-k policy denoising. |
| 188 | |
| 189 | Fuses per-row budget (3-regime k_schedule + linear decay, both keyed on |
| 190 | the full prompt KV length so the result is chunked-prefill invariant), |
| 191 | radix top-k threshold, and fixed retention (initial / window / diagonal). |
| 192 | |
| 193 | Args: |
| 194 | block_logits: Block-level OAM scores. |
| 195 | Shape: [num_batch, num_q_heads, max_Qb, max_Kb] |
| 196 | Dtype: bfloat16 (invalid positions set to -inf) |
| 197 | q_seq_lens: Q sequence length per request (current chunk). |
| 198 | Shape: [num_batch], Dtype: int32 |
| 199 | kv_seq_lens: KV sequence length per request (cumulative through current chunk). |
| 200 | Shape: [num_batch], Dtype: int32 |
| 201 | num_prompt_tokens: Full prompt KV-token count per request. For |
| 202 | chunked prefill pass the same value for every chunk of one |
| 203 | prompt; for normal prefill pass ``kv_seq_lens``. |
| 204 | Shape: [num_batch], Dtype: int32 |
| 205 | block_size: Stem sparse scoring block size (default 128). |
| 206 | alpha: Per-row budget decay factor (default 1.0 disables decay). |
| 207 | initial_blocks: Leading KV blocks always retained (default 4). |
| 208 | window_size: Recent diagonal-adjacent blocks always retained (default 4). |
| 209 | k_block_num_rate_medium: k_schedule multiplier when |
| 210 | 56 <= prompt_kv_blocks < 160 (default 0.2). |
| 211 | k_block_num_bias_medium: k_schedule bias in the medium regime (default 30). |
| 212 | k_block_num_rate_large: k_schedule multiplier when |
| 213 | prompt_kv_blocks >= 160 (default 0.1). |
| 214 | k_block_num_bias_large: k_schedule bias in the large regime (default 30). |
| 215 | |
| 216 | Returns: |
| 217 | mask: Per-block selection byte-mask. |
| 218 | Shape: [num_batch, num_q_heads, max_Qb, max_Kb] |
| 219 | Dtype: uint8 (1 = selected, 0 = skipped) |
| 220 | """ |
| 221 | return torch.ops.hpc.stem_tpd( |
| 222 | block_logits, |
| 223 | q_seq_lens, |
| 224 | kv_seq_lens, |
| 225 | num_prompt_tokens, |
| 226 | block_size, |
| 227 | alpha, |
| 228 | initial_blocks, |
| 229 | window_size, |
| 230 | k_block_num_rate_medium, |