From 50483a06dfeb33da126f38c7e2d99696cbdc998a Mon Sep 17 00:00:00 2001 From: whn09 Date: Fri, 23 Oct 2026 20:08:17 +0800 Subject: [PATCH] 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 --- deepinfer/environ.py | 4 ++++ deepinfer/model_executor/models/v4_modules/attn.py | 4 +++- .../model_executor/models/v4_modules/index_attn.py | 3 ++- deepinfer/worker/model_runner.py | 13 +++++++++++++ 4 files changed, 22 insertions(+), 2 deletions(-) diff --git a/deepinfer/environ.py b/deepinfer/environ.py index 3b71376..ccd6bb5 100644 --- a/deepinfer/environ.py +++ b/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). diff --git a/deepinfer/model_executor/models/v4_modules/attn.py b/deepinfer/model_executor/models/v4_modules/attn.py index 0fca8c8..7564ed0 100644 --- a/deepinfer/model_executor/models/v4_modules/attn.py +++ b/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) diff --git a/deepinfer/model_executor/models/v4_modules/index_attn.py b/deepinfer/model_executor/models/v4_modules/index_attn.py index 9d4f955..b3378ff 100644 --- a/deepinfer/model_executor/models/v4_modules/index_attn.py +++ b/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}), diff --git a/deepinfer/worker/model_runner.py b/deepinfer/worker/model_runner.py index dc8179a..742998a 100644 --- a/deepinfer/worker/model_runner.py +++ b/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)