{"task":"SVG-EAR","generated_at":"2026-10-19T20:11Z","runs":{"astra":{"run":"run_01M52C3P2ZA4NN3DYQGY18HC04","verdict":"FAIL","score":48.58633407346494,"submission":"Final submission","files":["svg/kmeans_utils.py","svg/models/hyvideo/attention.py","svg/models/wan/attention.py","svg/solver_extensions/adaptive.py","svg/solver_extensions/attention.py","svg/solver_extensions/mapping.py","svg/solver_extensions/policy.py","svg/solver_extensions/sparse.py"],"add":675,"del":5,"bytes":38077,"truncated":false,"patch":"diff --git a/svg/kmeans_utils.py b/svg/kmeans_utils.py\nindex a9049a8..1a97ad9 100644\n--- a/svg/kmeans_utils.py\n+++ b/svg/kmeans_utils.py\n@@ -917,7 +917,11 @@ def identify_dynamic_map_estimated(\n     min_kc_ratio=0.0,\n ):\n     \"\"\"Return a boolean block map with shape ``(B, H, Cq, Ck)``.\"\"\"\n-    raise NotImplementedError\n+    from .solver_extensions.mapping import estimated_map\n+\n+    if q_perm.shape[-1] != 128 or q_perm.dtype not in (torch.float16, torch.bfloat16):\n+        return identify_dynamic_map(qcentroids, kcentroids, qc_sz, kc_sz, top_p, min_kc_ratio)\n+    return estimated_map(q_perm, k_perm, kcentroids, qc_sz, kc_sz, top_p, min_kc_ratio)\n \n # --- Functions from analyze/dynamic_block_sparse_attention.py ---\n \n@@ -1025,7 +1029,9 @@ def scatter_mean_fused(src, size):\n         src: (cfg, num_heads, seq_len, dim) - Contiguous data\n         size: (cfg, num_heads, num_k_centroids) - Element count per label\n     \"\"\"\n-    raise NotImplementedError\n+    from .solver_extensions.attention import group_mean\n+\n+    return group_mean(src, size)\n \n @triton.jit\n def _dynamic_block_sparse_fwd_kernel(\n@@ -1454,7 +1460,12 @@ def dynamic_block_sparse_prune_fwd_flashinfer(\n     is_cpu: bool = True,\n ):\n     \"\"\"Return attention output with the same shape and dtype as ``q``.\"\"\"\n-    raise NotImplementedError\n+    from .solver_extensions.sparse import sparse_with_tail\n+\n+    return sparse_with_tail(\n+        q, k, v, k_centroids, v_centroids,\n+        block_mask_map.to(q.device), block_row_sz.to(q.device), block_col_sz.to(q.device),\n+    )\n \n # ---------------- Batch wrapper for cuVS KMeans -----------------\n \ndiff --git a/svg/models/hyvideo/attention.py b/svg/models/hyvideo/attention.py\nindex d5e482a..2c50015 100644\n--- a/svg/models/hyvideo/attention.py\n+++ b/svg/models/hyvideo/attention.py\n@@ -812,7 +812,88 @@ class Hunyuan_EARAttn_Processor(Hunyuan_SAPAttn_Processor2_0):\n     \"\"\"EAR processor hook.\"\"\"\n \n     def __init__(self, layer_idx):\n-        raise NotImplementedError\n+        super().__init__(layer_idx)\n+        self.prompt_length = int(self.prompt_length)\n+        self.top_p_kmeans = 0.95\n+        self.density_cap = 0.20\n+        self.adaptive_branches = 1\n+        self.kmeans_iter_init = 12\n+        self.kmeans_iter_step = 1\n+        self.zero_step_kmeans_init = False\n+        self.logging_file = None\n+        self.cluster_reuse = 2\n+\n+    def get_qk_norm(self, attn, query, key):\n+        from ...solver_extensions.attention import reference_rms_norm\n+\n+        return reference_rms_norm(query, attn.norm_q), reference_rms_norm(key, attn.norm_k)\n+\n+    def get_cu_max_seqlen(self, attention_mask, device):\n+        valid_length = int(attention_mask.sum().item())\n+        length = attention_mask.numel()\n+        self.prompt_length = valid_length - self.num_frame * self.frame_size\n+        offsets = torch.tensor([0, valid_length, length], device=device, dtype=torch.int32)\n+        return offsets, offsets, length, length\n+\n+    def flashinfer_attention(self, query, key, value, cu_max_seqlens):\n+        from ...solver_extensions.attention import dense_attention\n+\n+        valid_length = self.num_frame * self.frame_size + self.prompt_length\n+        output = torch.empty_like(query)\n+        output[:, :, :valid_length] = dense_attention(\n+            query[:, :, :valid_length], key[:, :, :valid_length], value[:, :, :valid_length]\n+        )\n+        if valid_length < query.shape[2]:\n+            output[:, :, valid_length:] = F.scaled_dot_product_attention(\n+                query[:, :, valid_length:], key[:, :, valid_length:], value[:, :, valid_length:]\n+            )\n+        return output\n+\n+    def semantic_aware_permutation(self, query, key, value, timestep, layer_idx,\n+                                   key_context=None, value_context=None):\n+        from ...solver_extensions.policy import clustered_inputs\n+\n+        return clustered_inputs(self, query, key, value, layer_idx, key_context, value_context)\n+\n+    def attention_core_logic(self, query, key, value, timestep, layer_idx, cu_max_seqlens):\n+        from ...solver_extensions.adaptive import cached_dense\n+\n+        if self.layer_idx < self.first_layers_fp or timestep[0] > self.first_times_fp:\n+            return self.flashinfer_attention(query, key, value, cu_max_seqlens)\n+        if cached_dense(self):\n+            return self.flashinfer_attention(query, key, value, cu_max_seqlens)\n+        video_length = self.num_frame * self.frame_size\n+        valid_length = video_length + self.prompt_length\n+        q_perm, k_perm, v_perm, block_map, row_sizes, col_sizes, order = self.semantic_aware_permutation(\n+            query[:, :, :video_length].contiguous(),\n+            key[:, :, :video_length].contiguous(),\n+            value[:, :, :video_length].contiguous(), timestep, layer_idx,\n+            key[:, :, video_length:valid_length], value[:, :, video_length:valid_length],\n+        )\n+        q_perm = torch.cat((q_perm, query[:, :, video_length:valid_length]), dim=2)\n+        k_perm = torch.cat((k_perm, key[:, :, video_length:valid_length]), dim=2)\n+        v_perm = torch.cat((v_perm, value[:, :, video_length:valid_length]), dim=2)\n+        block_map = F.pad(block_map, (0, 1, 0, 1), value=True)\n+        row_sizes = F.pad(row_sizes, (0, 1), value=self.prompt_length)\n+        col_sizes = F.pad(col_sizes, (0, 1), value=self.prompt_length)\n+        key_means = scatter_mean_fused(k_perm, col_sizes)\n+        value_means = scatter_mean_fused(v_perm, col_sizes)\n+        from ...solver_extensions.sparse import sparse_with_tail\n+\n+        result = sparse_with_tail(\n+            q_perm, k_perm, v_perm, key_means, value_means,\n+            block_map, row_sizes, col_sizes, variance_factors=self._active_variance_factors,\n+        )\n+        from ...solver_extensions.attention import inverse_tokens\n+\n+        output = torch.empty_like(query.transpose(1, 2), memory_format=torch.contiguous_format).transpose(1, 2)\n+        inverse_tokens(result[:, :, :video_length], order, output[:, :, :video_length])\n+        output[:, :, video_length:valid_length] = result[:, :, video_length:]\n+        if valid_length < query.shape[2]:\n+            output[:, :, valid_length:] = F.scaled_dot_product_attention(\n+                query[:, :, valid_length:], key[:, :, valid_length:], value[:, :, valid_length:]\n+            )\n+        return output\n \n def flashinfer_varlen_func(q, k, v, cu_seqlens_q, cu_seqlens_kv, max_seqlen_q, max_seqlen_kv):\n     \"\"\"\ndiff --git a/svg/models/wan/attention.py b/svg/models/wan/attention.py\nindex ec09c6e..504ed9b 100644\n--- a/svg/models/wan/attention.py\n+++ b/svg/models/wan/attention.py\n@@ -40,6 +40,8 @@ from .utils import (\n     generate_temporal_head_mask_mod,\n )\n \n+torch.cuda.memory._set_allocator_settings(\"expandable_segments:True\")\n+\n try:\n     # raise ImportError  # TODO: Remove this\n     sys.path.append(\"svg/kernels/build/\")\n@@ -564,4 +566,49 @@ class WanAttn_EARAttn_Processor(WanAttn_SAPAttn_Processor):\n     \"\"\"EAR processor hook.\"\"\"\n \n     def __init__(self, layer_idx):\n-        raise NotImplementedError\n+        super().__init__(layer_idx)\n+        self.top_p_kmeans = 0.95\n+        self.density_cap = 0.25\n+        self.adaptive_branches = 2\n+        self.kmeans_iter_init = 12\n+        self.kmeans_iter_step = 1\n+        self.zero_step_kmeans_init = False\n+        self.logging_file = None\n+        self.cluster_reuse = 2\n+\n+    def get_qk_norm(self, attn, query, key):\n+        from ...solver_extensions.attention import reference_rms_norm\n+\n+        return reference_rms_norm(query, attn.norm_q), reference_rms_norm(key, attn.norm_k)\n+\n+    def flash_attention(self, query, key, value):\n+        from ...solver_extensions.attention import dense_attention\n+\n+        return dense_attention(query, key, value)\n+\n+    def semantic_aware_permutation(self, query, key, value):\n+        from ...solver_extensions.policy import clustered_inputs\n+\n+        return clustered_inputs(self, query, key, value, self.layer_idx)\n+\n+    def attention_core_logic(self, query, key, value, timestep):\n+        from ...solver_extensions.adaptive import cached_dense\n+\n+        if self.layer_idx < self.first_layers_fp or timestep[0] > self.first_times_fp:\n+            return self.flash_attention(query, key, value)\n+        if cached_dense(self):\n+            return self.flash_attention(query, key, value)\n+        q_perm, k_perm, v_perm, block_map, row_sizes, col_sizes, order = self.semantic_aware_permutation(\n+            query, key, value\n+        )\n+        key_means = self._active_key_means\n+        value_means = scatter_mean_fused(v_perm, col_sizes)\n+        from ...solver_extensions.sparse import sparse_with_tail\n+\n+        output = sparse_with_tail(\n+            q_perm, k_perm, v_perm, key_means, value_means,\n+            block_map, row_sizes, col_sizes, variance_factors=self._active_variance_factors,\n+        )\n+        from ...solver_extensions.attention import inverse_tokens\n+\n+        return inverse_tokens(output, order)\ndiff --git a/svg/solver_extensions/adaptive.py b/svg/solver_extensions/adaptive.py\nnew file mode 100644\nindex 0000000..716b230\n--- /dev/null\n+++ b/svg/solver_extensions/adaptive.py\n@@ -0,0 +1,116 @@\n+import torch\n+\n+from .attention import group_mean\n+from .mapping import sampled_probabilities\n+from .sparse import _group_variance\n+\n+\n+def cached_dense(processor):\n+    calls = getattr(processor, \"_adaptive_calls\", 0)\n+    if (calls // processor.adaptive_branches) % 4 == 0:\n+        return False\n+    cached = getattr(processor, \"_adaptive_caps\", None)\n+    if cached is None:\n+        return False\n+    policy = cached[calls % processor.adaptive_branches]\n+    if policy is None or not bool((policy[0] >= 1.).all()):\n+        return False\n+    processor._adaptive_calls = calls + 1\n+    return True\n+\n+\n+def adaptive_map(processor, query, key, value, key_means, row_sizes, col_sizes,\n+                 threshold=.0005, probes=32, key_context=None, value_context=None):\n+    batch, heads, length, dim = query.shape\n+    calls = getattr(processor, \"_adaptive_calls\", 0)\n+    cached = getattr(processor, \"_adaptive_caps\", [None] * processor.adaptive_branches)\n+    maps = getattr(processor, \"_adaptive_maps\", [None] * processor.adaptive_branches)\n+    branch = calls % processor.adaptive_branches\n+    cached_policy = cached[branch]\n+    processor._adaptive_calls = calls + 1\n+    if cached_policy is not None and maps[branch] is not None and (calls // processor.adaptive_branches) % processor.cluster_reuse:\n+        processor._active_variance_factors = cached_policy[1]\n+        return maps[branch]\n+    key_vars = torch.empty_like(key_means)\n+    key_ends = col_sizes.cumsum(-1)\n+    _group_variance[(col_sizes.shape[-1], batch * heads)](\n+        key, key_means, key_ends, key_vars, length, col_sizes.shape[-1], dim, 32\n+    )\n+    probabilities = sampled_probabilities(query, key_means, key_vars, row_sizes, col_sizes, 128)\n+    sorted_probabilities, order = probabilities.sort(-1, descending=True)\n+    probability_mass = sorted_probabilities.cumsum(-1) - sorted_probabilities\n+    counts = col_sizes.unsqueeze(-2).expand_as(probabilities).gather(-1, order).cumsum(-1)\n+    if cached_policy is not None and (calls // processor.adaptive_branches) % 4:\n+        caps, factors = cached_policy\n+        processor._active_variance_factors = factors\n+        maps[branch] = _map_from_caps(probability_mass, counts, order, caps, length)\n+        processor._adaptive_maps = maps\n+        return maps[branch]\n+    value_means = group_mean(value, col_sizes)\n+    positions = ((torch.arange(probes, device=query.device) + .5) * length / probes).long()\n+    sampled_queries = query[:, :, positions].float()\n+    key_float, value_float = key.float(), value.float()\n+    scores = sampled_queries @ key_float.transpose(-2, -1) / dim**.5\n+    maximum = scores.max(-1, keepdim=True).values\n+    if key_context is not None and key_context.shape[2]:\n+        context_scores = sampled_queries @ key_context.float().transpose(-2, -1) / dim**.5\n+        maximum = torch.maximum(maximum, context_scores.max(-1, keepdim=True).values)\n+        context_weights = (context_scores - maximum).exp()\n+        context_numerator = context_weights @ value_context.float()\n+        context_denominator = context_weights.sum(-1, keepdim=True)\n+    else:\n+        context_numerator = 0.\n+        context_denominator = 0.\n+    weights = (scores - maximum).exp()\n+    reference = (weights @ value_float + context_numerator) / (weights.sum(-1, keepdim=True) + context_denominator)\n+    reference_norm = reference.square().mean()\n+    group_scores = sampled_queries @ key_means.float().transpose(-2, -1) / dim**.5\n+    variance_scores = sampled_queries.square() @ key_vars.float().transpose(-2, -1) / (2 * dim)\n+    group_scores += col_sizes.clamp_min(1).float().log().unsqueeze(-2)\n+    tail_candidates = []\n+    for factor in (0., .5, 1., 1.5):\n+        group_weights = (group_scores + factor * variance_scores - maximum).exp()\n+        group_weights.masked_fill_(col_sizes.unsqueeze(-2) == 0, 0.)\n+        tail_candidates.append(group_weights)\n+    query_groups = torch.searchsorted(row_sizes.cumsum(-1), positions.expand(batch, heads, -1).contiguous(), right=True)\n+    key_positions = torch.arange(length, device=query.device).expand(batch, heads, -1).contiguous()\n+    key_groups = torch.searchsorted(key_ends, key_positions, right=True)\n+    caps = torch.ones((batch, heads), device=query.device, dtype=torch.float32)\n+    factors = torch.ones_like(caps)\n+    remaining = torch.ones((batch, heads), device=query.device, dtype=torch.bool)\n+    for cap in (.12, .25, .50):\n+        remove = (probability_mass > .98) | (counts > length * cap)\n+        remove[..., :max(1, col_sizes.shape[-1] // 10)] = False\n+        candidate = torch.zeros_like(remove).scatter_(-1, order, ~remove)\n+        selected_groups = candidate.gather(-2, query_groups[..., None].expand(-1, -1, -1, col_sizes.shape[-1]))\n+        selected_tokens = selected_groups.gather(-1, key_groups.unsqueeze(-2).expand(-1, -1, probes, -1))\n+        selected_weights = weights * selected_tokens\n+        selected_numerator = selected_weights @ value_float + context_numerator\n+        selected_denominator = selected_weights.sum(-1, keepdim=True) + context_denominator\n+        best_error = torch.full_like(caps, float('inf'))\n+        best_factor = torch.ones_like(caps)\n+        for factor, group_weights in zip((0., .5, 1., 1.5), tail_candidates):\n+            tail_weights = group_weights * ~selected_groups\n+            numerator = selected_numerator + tail_weights @ value_means.float()\n+            denominator = selected_denominator + tail_weights.sum(-1, keepdim=True)\n+            candidate_output = numerator / denominator\n+            error = (candidate_output - reference).square().mean((-2, -1))\n+            best_factor = torch.where(error < best_error, factor, best_factor)\n+            best_error = torch.minimum(error, best_error)\n+        accept = remaining & (best_error <= threshold * reference_norm)\n+        caps = torch.where(accept, cap, caps)\n+        factors = torch.where(accept, best_factor, factors)\n+        remaining &= ~accept\n+    cached[branch] = caps, factors\n+    processor._adaptive_caps = cached\n+    processor._active_variance_factors = factors\n+    maps[branch] = _map_from_caps(probability_mass, counts, order, caps, length)\n+    processor._adaptive_maps = maps\n+    return maps[branch]\n+\n+\n+def _map_from_caps(probability_mass, counts, order, caps, length):\n+    remove = (probability_mass > .98) | (counts > length * caps[..., None, None])\n+    remove[..., :max(1, counts.shape[-1] // 10)] = False\n+    result = torch.zeros_like(remove).scatter_(-1, order, ~remove)\n+    return result | (caps[..., None, None] >= 1.)\ndiff --git a/svg/solver_extensions/attention.py b/svg/solver_extensions/attention.py\nnew file mode 100644\nindex 0000000..b39e56e\n--- /dev/null\n+++ b/svg/solver_extensions/attention.py\n@@ -0,0 +1,110 @@\n+import flashinfer\n+import torch\n+import triton\n+import triton.language as tl\n+from diffusers.models.normalization import RMSNorm\n+\n+\n+@triton.jit\n+def _reference_norm(Source, Weight, Output, ROWS: tl.constexpr, DIM: tl.constexpr,\n+                    PAD_DIM: tl.constexpr, EPS: tl.constexpr, BLOCK: tl.constexpr):\n+    rows = tl.program_id(0) * BLOCK + tl.arange(0, BLOCK)\n+    columns = tl.arange(0, PAD_DIM)\n+    mask = (rows[:, None] < ROWS) & (columns[None, :] < DIM)\n+    source = tl.load(Source + rows[:, None] * DIM + columns[None, :], mask, 0).to(tl.float32)\n+    scale = tl.rsqrt(tl.sum(source * source, 1) / DIM + EPS)\n+    normalized = (source * scale[:, None]).to(Output.dtype.element_ty).to(tl.float32)\n+    weight = tl.load(Weight + columns, columns < DIM, 0).to(tl.float32)\n+    output = normalized * weight[None, :]\n+    tl.store(Output + rows[:, None] * DIM + columns[None, :], output, mask)\n+\n+\n+def reference_rms_norm(source, norm):\n+    if norm is None:\n+        return source\n+    if (not isinstance(norm, RMSNorm) or norm.weight is None or getattr(norm, \"bias\", None) is not None\n+            or source.dtype not in (torch.float16, torch.bfloat16) or source.dtype != norm.weight.dtype\n+            or not source.is_cuda):\n+        return norm(source)\n+    source = source.contiguous()\n+    output = torch.empty_like(source)\n+    dim = source.shape[-1]\n+    rows = source.numel() // dim\n+    block = 16 if dim <= 256 else 1\n+    _reference_norm[(triton.cdiv(rows, block),)](\n+        source, norm.weight, output, rows, dim, triton.next_power_of_2(dim), norm.eps, block,\n+        num_warps=4 if dim <= 256 else 8,\n+    )\n+    return output\n+\n+\n+def dense_attention(query, key, value):\n+    outputs = []\n+    for batch_idx in range(query.shape[0]):\n+        output = flashinfer.single_prefill_with_kv_cache(\n+            query[batch_idx].transpose(0, 1),\n+            key[batch_idx].transpose(0, 1),\n+            value[batch_idx].transpose(0, 1),\n+            causal=False,\n+            backend=\"fa3\",\n+        )\n+        outputs.append(output)\n+    return torch.stack(outputs).transpose(1, 2)\n+\n+\n+@triton.jit\n+def _inverse_tokens(Source, Order, Output, LENGTH: tl.constexpr, HEADS: tl.constexpr,\n+                    DIM: tl.constexpr, SOURCE_BATCH: tl.constexpr, SOURCE_HEAD: tl.constexpr,\n+                    SOURCE_TOKEN: tl.constexpr, OUT_BATCH: tl.constexpr, OUT_HEAD: tl.constexpr,\n+                    OUT_TOKEN: tl.constexpr, BLOCK: tl.constexpr):\n+    head_index = tl.program_id(0)\n+    batch = head_index // HEADS\n+    head = head_index % HEADS\n+    tokens = tl.program_id(1) * BLOCK + tl.arange(0, BLOCK)\n+    dims = tl.arange(0, DIM)\n+    destinations = tl.load(Order + head_index * LENGTH + tokens, tokens < LENGTH, 0)\n+    source = Source + batch * SOURCE_BATCH + head * SOURCE_HEAD + tokens[:, None] * SOURCE_TOKEN + dims[None, :]\n+    output = Output + batch * OUT_BATCH + head * OUT_HEAD + destinations[:, None] * OUT_TOKEN + dims[None, :]\n+    values = tl.load(source, tokens[:, None] < LENGTH, 0)\n+    tl.store(output, values, tokens[:, None] < LENGTH)\n+\n+\n+def inverse_tokens(source, order, output=None):\n+    batch, heads, length, dim = source.shape\n+    if output is None:\n+        output = torch.empty((batch, length, heads, dim), device=source.device, dtype=source.dtype).transpose(1, 2)\n+    _inverse_tokens[(batch * heads, triton.cdiv(length, 64))](\n+        source, order, output, length, heads, dim, *source.stride()[:3], *output.stride()[:3], 64,\n+        num_warps=4,\n+    )\n+    return output\n+\n+\n+@triton.jit\n+def _group_mean(Src, Ends, Out, LENGTH: tl.constexpr, GROUPS: tl.constexpr,\n+                DIM: tl.constexpr, PAD_DIM: tl.constexpr, BLOCK: tl.constexpr):\n+    group = tl.program_id(0)\n+    head = tl.program_id(1)\n+    end = tl.load(Ends + head * GROUPS + group)\n+    start = tl.load(Ends + head * GROUPS + group - 1, group > 0, 0)\n+    dims = tl.arange(0, PAD_DIM)\n+    rows = tl.arange(0, BLOCK)\n+    accum = tl.full((BLOCK, PAD_DIM), 0, tl.float32)\n+    for offset in range(start, end, BLOCK):\n+        indices = offset + rows\n+        values = tl.load(Src + (head * LENGTH + indices[:, None]) * DIM + dims[None, :],\n+                         (indices[:, None] < end) & (dims[None, :] < DIM), 0)\n+        accum += values.to(tl.float32)\n+    mean = tl.sum(accum, 0) / tl.maximum(end - start, 1)\n+    tl.store(Out + (head * GROUPS + group) * DIM + dims, mean, dims < DIM)\n+\n+\n+def group_mean(src, size):\n+    batch, heads, length, dim = src.shape\n+    groups = size.shape[-1]\n+    ends = size.to(device=src.device, dtype=torch.int32).cumsum(-1, dtype=torch.int32)\n+    output = torch.empty((batch, heads, groups, dim), device=src.device, dtype=src.dtype)\n+    _group_mean[(groups, batch * heads)](\n+        src, ends, output, length, groups, dim, triton.next_power_of_2(dim), 32\n+    )\n+    return output\ndiff --git a/svg/solver_extensions/mapping.py b/svg/solver_extensions/mapping.py\nnew file mode 100644\nindex 0000000..1b2ff3c\n--- /dev/null\n+++ b/svg/solver_extensions/mapping.py\n@@ -0,0 +1,85 @@\n+import math\n+\n+import torch\n+import triton\n+import triton.language as tl\n+\n+\n+@triton.jit\n+def _sample_probabilities(Query, KeyMeans, KeyVars, RowEnds, ColSizes, Out,\n+                         LENGTH: tl.constexpr, DIM: tl.constexpr,\n+                         NUM_Q: tl.constexpr, NUM_K: tl.constexpr,\n+                         SCALE: tl.constexpr, VAR_SCALE: tl.constexpr,\n+                         SAMPLES: tl.constexpr, BLOCK_K: tl.constexpr):\n+    group = tl.program_id(0)\n+    head = tl.program_id(1)\n+    end = tl.load(RowEnds + head * NUM_Q + group)\n+    start = tl.load(RowEnds + head * NUM_Q + group - 1, group > 0, 0)\n+    sample = tl.arange(0, SAMPLES)\n+    rows = start + ((2 * sample + 1) * (end - start)) // (2 * SAMPLES)\n+    rows = tl.minimum(rows, LENGTH - 1)\n+    dims = tl.arange(0, DIM)\n+    query = tl.load(Query + (head * LENGTH + rows[:, None]) * DIM + dims[None, :])\n+    squared_query = (query.to(tl.float32) * query.to(tl.float32)).to(query.dtype)\n+    maximum = tl.full((SAMPLES,), -1e6, tl.float32)\n+    denominator = tl.full((SAMPLES,), 0., tl.float32)\n+    for offset in range(0, NUM_K, BLOCK_K):\n+        columns = offset + tl.arange(0, BLOCK_K)\n+        keys = tl.load(KeyMeans + (head * NUM_K + columns[None, :]) * DIM + dims[:, None], columns[None, :] < NUM_K, 0)\n+        variances = tl.load(KeyVars + (head * NUM_K + columns[None, :]) * DIM + dims[:, None], columns[None, :] < NUM_K, 0)\n+        sizes = tl.load(ColSizes + head * NUM_K + columns, columns < NUM_K, 0)\n+        scores = tl.dot(query, keys) * SCALE + tl.dot(squared_query, variances) * VAR_SCALE\n+        scores += tl.log2(tl.maximum(sizes, 1).to(tl.float32))[None, :]\n+        scores = tl.where(sizes[None, :] > 0, scores, -1e6)\n+        updated = tl.maximum(maximum, tl.max(scores, 1))\n+        weights = tl.exp2(scores - updated[:, None])\n+        denominator = denominator * tl.exp2(maximum - updated) + tl.sum(weights, 1)\n+        maximum = updated\n+    for offset in range(0, NUM_K, BLOCK_K):\n+        columns = offset + tl.arange(0, BLOCK_K)\n+        keys = tl.load(KeyMeans + (head * NUM_K + columns[None, :]) * DIM + dims[:, None], columns[None, :] < NUM_K, 0)\n+        variances = tl.load(KeyVars + (head * NUM_K + columns[None, :]) * DIM + dims[:, None], columns[None, :] < NUM_K, 0)\n+        sizes = tl.load(ColSizes + head * NUM_K + columns, columns < NUM_K, 0)\n+        scores = tl.dot(query, keys) * SCALE + tl.dot(squared_query, variances) * VAR_SCALE\n+        scores += tl.log2(tl.maximum(sizes, 1).to(tl.float32))[None, :]\n+        probabilities = tl.exp2(scores - maximum[:, None]) / denominator[:, None]\n+        probabilities = tl.where(sizes[None, :] > 0, probabilities, 0.)\n+        mean = tl.sum(probabilities, 0) / SAMPLES\n+        tl.store(Out + (head * NUM_Q + group) * NUM_K + columns, mean, columns < NUM_K)\n+\n+\n+def sampled_probabilities(query, key_means, key_vars, row_sizes, col_sizes, samples=64):\n+    batch, heads, length, dim = query.shape\n+    q_count, k_count = row_sizes.shape[-1], col_sizes.shape[-1]\n+    probabilities = torch.empty((batch, heads, q_count, k_count), device=query.device, dtype=torch.float32)\n+    _sample_probabilities[(q_count, batch * heads)](\n+        query, key_means, key_vars, row_sizes.cumsum(-1), col_sizes, probabilities,\n+        length, dim, q_count, k_count, math.log2(math.e) / dim**.5,\n+        math.log2(math.e) / (2 * dim), samples, 64 if samples >= 128 else 128,\n+        num_warps=4, num_stages=1,\n+    )\n+    return probabilities\n+\n+\n+def sampled_map(query, key_means, key_vars, row_sizes, col_sizes, top_p, min_ratio=.1, samples=64, density_cap=1.):\n+    probabilities = sampled_probabilities(query, key_means, key_vars, row_sizes, col_sizes, samples)\n+    sorted_probabilities, order = probabilities.sort(-1, descending=True)\n+    remove = sorted_probabilities.cumsum(-1) > top_p\n+    remove[..., 1:] = remove[..., :-1].clone()\n+    if density_cap < 1.:\n+        counts = col_sizes.unsqueeze(-2).expand_as(probabilities).gather(-1, order).cumsum(-1)\n+        remove |= counts > query.shape[-2] * density_cap\n+    remove[..., :max(1, int(min_ratio * col_sizes.shape[-1]))] = False\n+    return torch.zeros_like(remove).scatter_(-1, order, ~remove)\n+\n+\n+def estimated_map(query, key, key_means, row_sizes, col_sizes, top_p, min_ratio=0., samples=128, density_cap=1.):\n+    from .sparse import _group_variance\n+\n+    batch, heads, length, dim = query.shape\n+    key_vars = torch.empty_like(key_means)\n+    _group_variance[(col_sizes.shape[-1], batch * heads)](\n+        key, key_means, col_sizes.cumsum(-1), key_vars,\n+        length, col_sizes.shape[-1], dim, 32,\n+    )\n+    return sampled_map(query, key_means, key_vars, row_sizes, col_sizes, top_p, min_ratio, samples, density_cap)\ndiff --git a/svg/solver_extensions/policy.py b/svg/solver_extensions/policy.py\nnew file mode 100644\nindex 0000000..3fa702d\n--- /dev/null\n+++ b/svg/solver_extensions/policy.py\n@@ -0,0 +1,42 @@\n+from ..kernels.triton.permute import permute_tensor_by_labels_triton\n+from .adaptive import adaptive_map\n+from .attention import group_mean\n+\n+\n+def clustered_inputs(processor, query, key, value, layer_idx, key_context=None, value_context=None):\n+    batch, heads, length, dim = query.shape\n+    calls = getattr(processor, \"_adaptive_calls\", 0)\n+    branch = calls % processor.adaptive_branches\n+    layouts = getattr(processor, \"_cluster_layouts\", [None] * processor.adaptive_branches)\n+    layout = layouts[branch]\n+    if layout is None or (calls // processor.adaptive_branches) % processor.cluster_reuse == 0:\n+        if processor.adaptive_branches > 1:\n+            centers = getattr(processor, \"_branch_centers\", [None] * processor.adaptive_branches)\n+            previous = centers[branch]\n+            processor.centroids_init = previous is not None\n+            if previous is not None:\n+                processor.q_centroids, processor.k_centroids = previous\n+        q_labels, q_means, q_sizes, _, k_labels, k_means, k_sizes, _ = processor.kmeans_clustering(\n+            query, key, layer_idx\n+        )\n+        if processor.adaptive_branches > 1:\n+            centers[branch] = q_means, k_means\n+            processor._branch_centers = centers\n+        q_perm, q_order = permute_tensor_by_labels_triton(query, q_labels, dim=2)\n+        k_perm, k_order = permute_tensor_by_labels_triton(key, k_labels, dim=2)\n+        q_sizes = q_sizes.reshape(batch, heads, -1)\n+        k_sizes = k_sizes.reshape(batch, heads, -1)\n+        layouts[branch] = q_order, k_order, q_sizes, k_sizes\n+        processor._cluster_layouts = layouts\n+    else:\n+        q_order, k_order, q_sizes, k_sizes = layout\n+        q_perm, _ = permute_tensor_by_labels_triton(query, None, dim=2, sorted_indices=q_order)\n+        k_perm, _ = permute_tensor_by_labels_triton(key, None, dim=2, sorted_indices=k_order)\n+    v_perm, _ = permute_tensor_by_labels_triton(value, None, dim=2, sorted_indices=k_order)\n+    k_means = group_mean(k_perm, k_sizes)\n+    processor._active_key_means = k_means\n+    block_map = adaptive_map(\n+        processor, q_perm, k_perm, v_perm, k_means, q_sizes, k_sizes,\n+        key_context=key_context, value_context=value_context,\n+    )\n+    return q_perm, k_perm, v_perm, block_map, q_sizes, k_sizes, q_order\ndiff --git a/svg/solver_extensions/sparse.py b/svg/solver_extensions/sparse.py\nnew file mode 100644\nindex 0000000..5f74a15\n--- /dev/null\n+++ b/svg/solver_extensions/sparse.py\n@@ -0,0 +1,178 @@\n+import math\n+\n+import flashinfer\n+import torch\n+import triton\n+import triton.language as tl\n+\n+from ..flashinfer_patch import flashinfer_patch_enabled\n+\n+\n+_workspaces = {}\n+\n+\n+@triton.jit\n+def _group_variance(Src, Means, Ends, Out, LENGTH: tl.constexpr,\n+                    GROUPS: tl.constexpr, DIM: tl.constexpr, BLOCK: tl.constexpr):\n+    group = tl.program_id(0)\n+    head = tl.program_id(1)\n+    end = tl.load(Ends + head * GROUPS + group)\n+    start = tl.load(Ends + head * GROUPS + group - 1, group > 0, 0)\n+    dims = tl.arange(0, DIM)\n+    rows = tl.arange(0, BLOCK)\n+    mean = tl.load(Means + (head * GROUPS + group) * DIM + dims).to(tl.float32)\n+    accum = tl.full((BLOCK, DIM), 0, tl.float32)\n+    for offset in range(start, end, BLOCK):\n+        indices = offset + rows\n+        values = tl.load(Src + (head * LENGTH + indices[:, None]) * DIM + dims[None, :], indices[:, None] < end, 0).to(tl.float32)\n+        residual = tl.where(indices[:, None] < end, values - mean[None, :], 0.)\n+        accum += residual * residual\n+    variance = tl.sum(accum, 0) / tl.maximum(end - start, 1)\n+    tl.store(Out + (head * GROUPS + group) * DIM + dims, variance)\n+\n+\n+def exact_sparse(query, key, value, block_map, row_sizes, col_sizes):\n+    batch, heads, length, dim = query.shape\n+    cache_key = (query.device, query.dtype)\n+    required = int((block_map * col_sizes.unsqueeze(-2)).sum().item())\n+    if cache_key not in _workspaces:\n+        workspace = torch.empty(512 * 1024 * 1024, device=query.device, dtype=torch.uint8)\n+        wrapper = flashinfer.sparse.VariableBlockSparseAttentionWrapper(workspace, backend='auto')\n+        capacity = 0\n+    else:\n+        wrapper, capacity = _workspaces[cache_key]\n+    if required >= capacity:\n+        capacity = max(1048576, triton.cdiv(required + 1, 1048576) * 1048576)\n+        indices = torch.empty(capacity, device=query.device, dtype=torch.int32)\n+        wrapper.reset_workspace_buffer(\n+            float_workspace_buffer=wrapper._float_workspace_buffer,\n+            int_workspace_buffer=wrapper._int_workspace_buffer,\n+            vector_sparse_indices_buffer=indices,\n+            vector_sparse_indptr_buffer=wrapper._vector_sparse_indptr_buffer,\n+        )\n+    _workspaces[cache_key] = wrapper, capacity\n+    with flashinfer_patch_enabled():\n+        wrapper.plan(\n+            block_mask_map=block_map.reshape(batch * heads, *block_map.shape[-2:]),\n+            block_row_sz=row_sizes.reshape(batch * heads, -1),\n+            block_col_sz=col_sizes.reshape(batch * heads, -1),\n+            num_qo_heads=batch * heads,\n+            num_kv_heads=batch * heads,\n+            head_dim=dim,\n+            q_data_type=query.dtype,\n+            kv_data_type=key.dtype,\n+        )\n+    output, lse = wrapper.run(\n+        query.reshape(batch * heads, length, dim),\n+        key.reshape(batch * heads, length, dim),\n+        value.reshape(batch * heads, length, dim),\n+        return_lse=True,\n+    )\n+    return output.reshape_as(query), lse\n+\n+\n+@triton.jit\n+def _tail_merge(Query, KeyMeans, ValueMeans, KeyVars, ColSizes, RowEnds, BlockMap,\n+                Exact, ExactLSE, Out, VarianceFactors,\n+                LENGTH: tl.constexpr, DIM: tl.constexpr, NUM_Q: tl.constexpr,\n+                NUM_K: tl.constexpr, SCALE: tl.constexpr,\n+                LSE_SCALE: tl.constexpr, VAR_SCALE: tl.constexpr, PER_HEAD_VAR: tl.constexpr, SEARCH: tl.constexpr,\n+                BLOCK_Q: tl.constexpr, BLOCK_K: tl.constexpr):\n+    head = tl.program_id(1)\n+    variance_scale = VAR_SCALE\n+    if PER_HEAD_VAR:\n+        variance_scale *= tl.load(VarianceFactors + head)\n+    rows = tl.program_id(0) * BLOCK_Q + tl.arange(0, BLOCK_Q)\n+    dims = tl.arange(0, DIM)\n+    low = tl.full((BLOCK_Q,), 0, tl.int32)\n+    high = tl.full((BLOCK_Q,), NUM_Q, tl.int32)\n+    for iteration in range(SEARCH):\n+        middle = (low + high) // 2\n+        end = tl.load(RowEnds + head * NUM_Q + middle, middle < NUM_Q, LENGTH)\n+        below = rows >= end\n+        low = tl.where(below, middle + 1, low)\n+        high = tl.where(below, high, middle)\n+    groups = tl.minimum(low, NUM_Q - 1)\n+    query = tl.load(Query + (head * LENGTH + rows[:, None]) * DIM + dims[None, :], rows[:, None] < LENGTH, 0)\n+    maximum = tl.full((BLOCK_Q,), -1.0e6, tl.float32)\n+    denominator = tl.full((BLOCK_Q,), 0, tl.float32)\n+    numerator = tl.full((BLOCK_Q, DIM), 0, tl.float32)\n+    for start in range(0, NUM_K, BLOCK_K):\n+        columns = start + tl.arange(0, BLOCK_K)\n+        keys = tl.load(KeyMeans + (head * NUM_K + columns[None, :]) * DIM + dims[:, None], columns[None, :] < NUM_K, 0)\n+        sizes = tl.load(ColSizes + head * NUM_K + columns, columns < NUM_K, 0)\n+        selected = tl.load(BlockMap + (head * NUM_Q + groups[:, None]) * NUM_K + columns[None, :], columns[None, :] < NUM_K, 1)\n+        scores = tl.dot(query, keys).to(tl.float32) * SCALE\n+        if VAR_SCALE > 0:\n+            variances = tl.load(KeyVars + (head * NUM_K + columns[None, :]) * DIM + dims[:, None], columns[None, :] < NUM_K, 0)\n+            squared_query = (query.to(tl.float32) * query.to(tl.float32)).to(query.dtype)\n+            scores += tl.dot(squared_query, variances) * variance_scale\n+        scores += tl.log2(tl.maximum(sizes, 1).to(tl.float32))[None, :]\n+        scores = tl.where((~selected) & (sizes[None, :] > 0), scores, -1.0e6)\n+        updated = tl.maximum(maximum, tl.max(scores, 1))\n+        weights = tl.exp2(scores - updated[:, None])\n+        weights = tl.where((~selected) & (sizes[None, :] > 0), weights, 0)\n+        alpha = tl.exp2(maximum - updated)\n+        denominator = denominator * alpha + tl.sum(weights, 1)\n+        numerator *= alpha[:, None]\n+        values = tl.load(ValueMeans + (head * NUM_K + columns[:, None]) * DIM + dims[None, :], columns[:, None] < NUM_K, 0)\n+        numerator += tl.dot(weights.to(values.dtype), values)\n+        maximum = updated\n+    tail_lse = maximum + tl.log2(denominator)\n+    exact_lse = tl.load(ExactLSE + head * LENGTH + rows, rows < LENGTH, 0) * LSE_SCALE\n+    tail_fraction = 1.0 / (1.0 + tl.exp2(exact_lse - tail_lse))\n+    exact = tl.load(Exact + (head * LENGTH + rows[:, None]) * DIM + dims[None, :], rows[:, None] < LENGTH, 0).to(tl.float32)\n+    approx = numerator / tl.maximum(denominator, 1.0e-20)[:, None]\n+    result = exact + tail_fraction[:, None] * (approx - exact)\n+    tl.store(Out + (head * LENGTH + rows[:, None]) * DIM + dims[None, :], result, rows[:, None] < LENGTH)\n+\n+\n+def sparse_with_tail(query, key, value, key_means, value_means, block_map, row_sizes, col_sizes, lse_scale=1.0, variance_factor=1.0, variance_factors=None, split_heads=True):\n+    batch, heads, length, dim = query.shape\n+    if dim not in (64, 128, 256) or query.dtype not in (torch.float16, torch.bfloat16):\n+        return torch.nn.functional.scaled_dot_product_attention(query, key, value)\n+    selected = (block_map * col_sizes.unsqueeze(-2)).sum(-1)\n+    densities = (selected * row_sizes).sum(-1).float() / (length * length)\n+    if split_heads:\n+        dense_indices = (densities.flatten() > .4).nonzero().flatten()\n+        sparse_indices = (densities.flatten() <= .4).nonzero().flatten()\n+        if dense_indices.numel() and sparse_indices.numel():\n+            from .attention import dense_attention\n+\n+            output = torch.empty_like(query).flatten(0, 1)\n+            dense_inputs = [source.flatten(0, 1).index_select(0, dense_indices).unsqueeze(0)\n+                            for source in (query, key, value)]\n+            output.index_copy_(0, dense_indices, dense_attention(*dense_inputs).squeeze(0))\n+            del dense_inputs\n+            sparse_inputs = [source.flatten(0, 1).index_select(0, sparse_indices).unsqueeze(0)\n+                             for source in (query, key, value, key_means, value_means, block_map, row_sizes, col_sizes)]\n+            sparse_factors = None if variance_factors is None else variance_factors.flatten().index_select(0, sparse_indices)\n+            sparse_output = sparse_with_tail(*sparse_inputs, lse_scale=lse_scale,\n+                                             variance_factor=variance_factor, variance_factors=sparse_factors,\n+                                             split_heads=False)\n+            output.index_copy_(0, sparse_indices, sparse_output.squeeze(0))\n+            return output.reshape_as(query)\n+    density = densities.mean().item()\n+    if density > 0.45:\n+        from .attention import dense_attention\n+\n+        return dense_attention(query, key, value)\n+    key_vars = key_means\n+    if variance_factor:\n+        key_vars = torch.empty_like(key_means)\n+        key_ends = col_sizes.cumsum(-1)\n+        _group_variance[(col_sizes.shape[-1], batch * heads)](\n+            key, key_means, key_ends, key_vars, length, col_sizes.shape[-1], dim, 32\n+        )\n+    output, lse = exact_sparse(query, key, value, block_map, row_sizes, col_sizes)\n+    row_ends = row_sizes.cumsum(-1).contiguous()\n+    batch, heads, length, dim = query.shape\n+    merged = torch.empty_like(query)\n+    _tail_merge[(triton.cdiv(length, 128), batch * heads)](\n+        query, key_means, value_means, key_vars, col_sizes, row_ends, block_map,\n+        output, lse, merged, variance_factors, length, dim, row_sizes.shape[-1], col_sizes.shape[-1],\n+        math.log2(math.e) / math.sqrt(dim), lse_scale, variance_factor * math.log2(math.e) / (2 * dim),\n+        variance_factors is not None,\n+        (row_sizes.shape[-1] + 1).bit_length(), 128, 64, num_warps=4, num_stages=1,\n+    )\n+    return merged\n"}}}