Neuron: context-length buckets for dense compressed-KV scopes and indexer logits (#29)

DEEPINFER_NEURON_CTX_BUCKETS: the model runner sets the step's context limit
to the smallest bucket that holds every sequence, and the dense kv_ca_index /
indexer logits are sized to it instead of max_model_len (one more set of
attention graphs a bucket). Prefill cost scaled with the limit: a 128-token
chunk ~9 s at 16384 against ~5.3 s at 4352. Host tiny check: same tokens as
without buckets up to near ties.

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>

deepinfer/environ.py

@@ -126,6 +126,10 @@ class envs:
     # with more routed rows reruns the layer's MoE densely. Off by default: measured slower (108 vs 58 ms a layer for
     # a 128-token chunk at EP=64, and every layer overflowed at 1024), issue #26.
     DEEPINFER_NEURON_MOE_CAPACITY = EnvVar.int(default=0)
+    # Neuron only, region graphs: context-length buckets (tokens, comma-separated, ascending). A step whose longest
+    # sequence fits a bucket sizes its dense compressed-KV scopes and indexer logits to the bucket instead of
+    # max_model_len (one more set of attention graphs a bucket; issue #29).
+    DEEPINFER_NEURON_CTX_BUCKETS = EnvVar.str(default="")
     # Neuron only: run the sampler's process_sampling as one compiled graph instead of ~1900 eager ops a step.
     DEEPINFER_NEURON_COMPILE_SAMPLER = EnvVar.bool(default=True)
     # Neuron only: compile the decode step with fullgraph=True (raise on the first graph break).

deepinfer/model_executor/models/v4_modules/attn.py

@@ -768,7 +768,9 @@ class V4Attention(V4AttentionBase):
         # into the runtime by the first YOCO consumer layer
         shared_hidden_bf16 = cross_layer_runtime.shared_hidden
         block_size = kv_cache.ca_cache.shape[1] * self.kv_compress_ratio if kv_cache.ca_cache is not None else None
-        input_metadata.compute_attn_info(self.max_model_len, self.window_size, block_size)
+        # Neuron: the step's context bucket (model_runner), so dense compressed scopes span the context, not the limit.
+        ctx_limit = getattr(input_metadata, "neuron_ctx_limit", None) or self.max_model_len
+        input_metadata.compute_attn_info(ctx_limit, self.window_size, block_size)
         use_cp = input_metadata.is_prompt and PG.context_parallel_group().size() > 1
         shared_hidden_bf16 = self._resolve_shared_hidden(hidden_states_bf16, shared_hidden_bf16)
         projected_kv, index_weights = self._project_current_hidden_inputs(hidden_states_bf16)

deepinfer/model_executor/models/v4_modules/index_attn.py

@@ -456,7 +456,8 @@ class MQA(nn.Module):
                 input_metadata.klen[self.ratio_key].unsqueeze(1),
                 input_metadata.block_tables,
                 schedule_metadata,
-                self.max_kv_len,
+                (getattr(input_metadata, "neuron_ctx_limit", None) or self.max_kv_len * self.kv_compress_ratio)
+                // self.kv_compress_ratio,
                 indices=input_metadata.query_indices,
                 # TODO(NPU): Setting `use_ascend_kv_layout` to `True` on NPU may be faster, but the KV cache layout will be changed.
                 **({} if PLATFORM.is_gpu else {"use_ascend_kv_layout": False}),

deepinfer/worker/model_runner.py

@@ -173,6 +173,17 @@ class ModelRunner:
             return self._prepare_input_driver(seq_metadata_list, metadata_zos, total_q)
         return self._prepare_input_follower()
 
+    def _neuron_ctx_limit(self, seq_metadata_list: list[SequenceMetadata]) -> int | None:
+        """The smallest DEEPINFER_NEURON_CTX_BUCKETS entry holding every sequence's context, None (max_model_len)
+        when none does."""
+        buckets = getattr(self, "_ctx_buckets", None)
+        if buckets is None:
+            buckets = self._ctx_buckets = sorted(int(b) for b in envs.DEEPINFER_NEURON_CTX_BUCKETS.split(",") if b.strip())
+        if not buckets:
+            return None
+        longest = max(seq.seqlen_k for seq in seq_metadata_list)
+        return next((b for b in buckets if longest <= b), None)
+
     @staticmethod
     def _graph_input_metadata(cuda_graph_used: CudaGraphData, num_reqs: int) -> InputMetadata:
         input_metadata = cuda_graph_used.input_metadata_device
@@ -212,6 +223,8 @@ class ModelRunner:
                 prompt_as_decode=prompt_as_decode,
             )
             input_metadata = self._graph_input_metadata(cuda_graph_used, num_reqs)
+            if PLATFORM.is_neuron:
+                input_metadata.neuron_ctx_limit = self._neuron_ctx_limit(seq_metadata_list)
             if prompt_as_decode:
                 # The graph runs on its decode-mode metadata; sampling and post-processing see a prompt step.
                 input_metadata = copy.copy(input_metadata)