From d43b7e25a6c27c800f02f9808b0a5104d3789eba Mon Sep 17 00:00:00 2001 From: gf239 Date: Wed, 9 Sep 2026 02:11:39 +0200 Subject: [PATCH 1/3] perf(mamba): trim per-call overhead on the GDN and causal_conv1d decode path Three local rewrites on lines that execute once per GDN layer per engine step -- 48 times a step on a Qwen3-Next-class hybrid. None changes a value, a dtype, a shape, an alias, or a launch; each removes Python work that the GPU is waiting on when cudagraphs are piecewise. qwen_gdn_linear_attn.py, four sites: `self.conv1d.weight` was read three times to build one view (once for the view, twice for its sizes). Each read is two nn.Module __getattr__ hops (_modules, then _parameters). Bind it once. conv1d is a plain ColumnParallelLinear whose weight is a registered Parameter with no property or hook in the way, so the three reads returned the same object and the hoist is exact. causal_conv1d.py, dtype guards: `x.to(conv_state.dtype)` on entry and `out.to(original_x_dtype)` on exit are dispatched unconditionally, but under the default mamba_cache_dtype="auto" the conv state has the model dtype and both are no-ops. Tensor.to on a matching dtype returns self, so guarding on dtype inequality yields the identical object and only skips the ATen dispatch. The `out = x` aliasing when no output buffer is passed is unchanged: x is the same object either way. causal_conv1d.py, next_power_of_2: replace triton.next_power_of_2 with vllm.utils.math_utils.next_power_of_2, which 18 other modules already use. The two agree for every n >= 1; state_len here is width-1 (or width-1 + seqlen-1 on the spec-decode path) with width >= 2, so n >= 1 always holds. The one input where they differ, n == 0, needs width == 1, which has no conv-state row and no kernel branch; stock rejected it at Triton compile time via tl.arange(0, 0). The helper is one shift instead of an eight-step bit twiddle, and the module no longer reaches through the Triton namespace for a scalar. Measured per call on the host (torch 2.14, CPU): 0.92 us, 0.31 us, 0.31 us and 0.19 us respectively -- about 85 us per step at 48 layers. Co-Authored-By: Claude Fable 5.1 Signed-off-by: gf239 --- .../layers/mamba/gdn/qwen_gdn_linear_attn.py | 20 ++++++++----------- .../layers/mamba/ops/causal_conv1d.py | 13 +++++++----- 2 files changed, 16 insertions(+), 17 deletions(-) diff --git a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py index bbbb1c8bd5ad..680492e16e94 100644 --- a/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py +++ b/vllm/model_executor/layers/mamba/gdn/qwen_gdn_linear_attn.py @@ -1331,9 +1331,8 @@ def _forward_core( a = a[:num_actual_tokens] # 1. Convolution sequence transformation - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) + w = self.conv1d.weight + conv_weights = w.view(w.size(0), w.size(2)) if spec_sequence_masks is not None: if attn_metadata.num_prefills == 0 and attn_metadata.num_decodes == 0: @@ -1609,9 +1608,8 @@ def _forward_core_decode_aiter( ssm_state = self_kv_cache[1] # 1. Convolution sequence transformation - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) + w = self.conv1d.weight + conv_weights = w.view(w.size(0), w.size(2)) mixed_qkv_non_spec, b, a = ( gdn_aiter_fused_reshape_causal_conv1d_update_single_token( @@ -1681,9 +1679,8 @@ def _forward_core_decode_non_spec( b = b[:num_actual_tokens] a = a[:num_actual_tokens] - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) + w = self.conv1d.weight + conv_weights = w.view(w.size(0), w.size(2)) mixed_qkv_non_spec = causal_conv1d_update( mixed_qkv, conv_state, @@ -1731,9 +1728,8 @@ def _forward_core_decode_spec_fused_norm( if is_conv_state_dim_first() else self.kv_cache[0].transpose(-1, -2) ) - conv_weights = self.conv1d.weight.view( - self.conv1d.weight.size(0), self.conv1d.weight.size(2) - ) + w = self.conv1d.weight + conv_weights = w.view(w.size(0), w.size(2)) mixed_qkv = causal_conv1d_update( mixed_qkv[:num_actual_tokens], conv_state, diff --git a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py index 8335c849666d..ab36112ebdc5 100644 --- a/vllm/model_executor/layers/mamba/ops/causal_conv1d.py +++ b/vllm/model_executor/layers/mamba/ops/causal_conv1d.py @@ -10,6 +10,7 @@ from vllm.platforms import current_platform from vllm.triton_utils import tl, triton +from vllm.utils.math_utils import next_power_of_2 from vllm.v1.attention.backends.utils import NULL_BLOCK_ID, PAD_SLOT_ID @@ -577,7 +578,7 @@ def causal_conv1d_fn( dim, cu_seqlen = x.shape _, width = weight.shape state_len = width - 1 - np2_statelen = triton.next_power_of_2(state_len) + np2_statelen = next_power_of_2(state_len) padded_batch = query_start_loc.size(0) - 1 stride_x_dim = x.stride(0) @@ -756,7 +757,7 @@ def grid(META): num_stages=2, launch_pdl=current_platform.is_arch_support_pdl(), ) - return out.to(original_x_dtype) + return out if out.dtype == original_x_dtype else out.to(original_x_dtype) @triton.jit(do_not_specialize_on_alignment=["num_cache_lines"]) @@ -1157,7 +1158,9 @@ def causal_conv1d_update( assert activation in ["silu", "swish"] original_x_dtype = x.dtype - x = x.to(conv_state.dtype) + conv_state_dtype = conv_state.dtype + if original_x_dtype != conv_state_dtype: + x = x.to(conv_state_dtype) if out is None: out = x else: @@ -1221,7 +1224,7 @@ def causal_conv1d_update( state_len = width - 1 + (seqlen - 1) # effective state_len needed else: state_len = width - 1 - np2_statelen = triton.next_power_of_2(state_len) + np2_statelen = next_power_of_2(state_len) def grid(META): return ( @@ -1276,7 +1279,7 @@ def grid(META): ) if unsqueeze: out = out.squeeze(-1) - return out.to(original_x_dtype) + return out if out.dtype == original_x_dtype else out.to(original_x_dtype) if current_platform.is_cpu(): From 1fd630125456b97df8b48b3aa881a09cae36788a Mon Sep 17 00:00:00 2001 From: gf239 Date: Wed, 9 Sep 2026 02:11:39 +0200 Subject: [PATCH 2/3] perf(spec_decode): drop a redundant sum, a no-op rebase and a double slice Three once-per-step rewrites in the V2 runner's speculative-decoding path. Values, dtypes and shapes are unchanged; what goes away is a second pass over the same data, one GPU kernel launch with its allocation, and a throwaway view. model_runner.py, input prep: total_num_draft_tokens was computed with a separate `.sum()` and total_num_logits derived from it, and then a few lines later `np.cumsum(num_logits, out=cu_num_logits_np[1:])` walked the same array again. The cumsum's last element is that total, so derive both from it after the cumsum. Nothing in between reads either name. With num_reqs == 0 the cumsum target is empty and cu_num_logits_np[-1] is the leading 0, which is what the sum produced. rejection_sampler.py, chunked verify: each chunk rebased its cumulative logit offsets with `- lo`. For the single-chunk case -- every step whose logit rows fit in one chunk, i.e. the common one -- and for the first chunk of a split batch, lo is 0 and the subtraction is a no-op that still costs an aten::sub dispatch, an allocation and a cudaLaunchKernel, on the GPU-idle gap between the target forward and the first draft. Guard on the value rather than on the invariant: when lo is 0 take the plain slices. Those are views of the input batch buffers; every consumer of chunk_cu_num_logits only reads them (the verify kernels tl.load, and the one escaping use is cloned), so the aliasing is not observable. speculator.py, draft metadata: seq_lens was pre-sliced to num_reqs_padded and then build_attn_metadata re-applied the identical bound as its first use of the buffer, so the call site built buf[:n][:n]. Pass the buffer as every other caller of build_attn_metadata already does; the callee's slice is the only one needed. Co-Authored-By: Claude Fable 5.1 Signed-off-by: gf239 --- vllm/v1/worker/gpu/model_runner.py | 6 ++++-- vllm/v1/worker/gpu/spec_decode/rejection_sampler.py | 12 ++++++++++-- vllm/v1/worker/gpu/spec_decode/speculator.py | 4 +++- 3 files changed, 17 insertions(+), 5 deletions(-) diff --git a/vllm/v1/worker/gpu/model_runner.py b/vllm/v1/worker/gpu/model_runner.py index 7dec923c0a8f..9aa75455005e 100644 --- a/vllm/v1/worker/gpu/model_runner.py +++ b/vllm/v1/worker/gpu/model_runner.py @@ -1243,8 +1243,6 @@ def prepare_inputs( count=num_reqs, ) num_bonus_tokens = self.model_state.num_new_sampled_tokens_per_step - total_num_draft_tokens = int(num_draft_tokens_per_req.sum()) - total_num_logits = num_reqs * num_bonus_tokens + total_num_draft_tokens num_logits = num_draft_tokens_per_req + num_bonus_tokens # combine_sampled_and_draft_tokens places a request's logits rows # at [query_end - num_logits, query_end). Fewer query rows than @@ -1253,6 +1251,10 @@ def prepare_inputs( cu_num_logits_np = np.empty(num_reqs + 1, dtype=np.int32) cu_num_logits_np[0] = 0 np.cumsum(num_logits, out=cu_num_logits_np[1:]) + # The cumsum's last element IS the total, so the separate .sum() + # above was a second pass over the same data. + total_num_logits = int(cu_num_logits_np[-1]) + total_num_draft_tokens = total_num_logits - num_reqs * num_bonus_tokens cu_num_logits = async_copy_to_gpu(cu_num_logits_np, device=self.device) adaptive_verification = ( diff --git a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py index a9027e0550f8..c555c1a859b5 100644 --- a/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py +++ b/vllm/v1/worker/gpu/spec_decode/rejection_sampler.py @@ -211,8 +211,16 @@ def _verify_in_chunks( for start, end in request_chunks: lo = int(cu_num_logits_np[start]) hi = int(cu_num_logits_np[end]) - chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo - chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] - lo + if lo: + chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] - lo + chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] - lo + else: + # cu_num_logits always starts at 0, so the single-chunk case + # (and the first chunk of a split batch) rebases by nothing. + # Skip the GPU sub: a kernel launch and an allocation for a no-op. + # NOTE: these are read-only views of the input batch buffers. + chunk_cu_num_logits_np = cu_num_logits_np[start : end + 1] + chunk_cu_num_logits = input_batch.cu_num_logits[start : end + 1] # draft_logits uses persistent request-state indices and stays global. processed_logits, sampled, num_sampled = self._verify( logits[lo:hi], diff --git a/vllm/v1/worker/gpu/spec_decode/speculator.py b/vllm/v1/worker/gpu/spec_decode/speculator.py index ecf36ce517fd..6a278a63c536 100644 --- a/vllm/v1/worker/gpu/spec_decode/speculator.py +++ b/vllm/v1/worker/gpu/spec_decode/speculator.py @@ -315,7 +315,9 @@ def _build_draft_attn_metadata( ], query_start_loc_cpu=query_start_loc_cpu, max_query_len=max_query_len, - seq_lens=self.input_buffers.seq_lens[:num_reqs_padded], + # build_attn_metadata re-applies this exact bound as its first + # statement, so pre-slicing here only builds a throwaway view. + seq_lens=self.input_buffers.seq_lens, dcp_local_seq_lens=( None if dcp_local_seq_lens is None From 204bf8cd2cf9c264bb3166ca504076b25e11953e Mon Sep 17 00:00:00 2001 From: gf239 Date: Wed, 9 Sep 2026 02:11:39 +0200 Subject: [PATCH 3/3] perf(v1/attn): skip KV cache groups a drafter does not own when building metadata build_attn_metadata loops over every KV cache group and, for each, selects the group's block table and slot mapping, resolves the causal flag, builds a CommonAttentionMetadata (27 fields) and only then iterates that group's attention groups. A drafter initialises its attention groups from active_layer_names alone, so init_attn_backend appends an empty list for every group the draft model has no layer in. On a hybrid target with 16 full-attention and 48 linear-attention layers there are four KV cache groups and the drafter owns one; the other three were constructed and discarded on both draft metadata builds of every step. Test the group list first and continue when it is empty. The loop body writes nothing outside the inner loop -- the sole write into the returned mapping is per attention group -- so the result is byte-identical. The cost when every group is populated is one truthiness test per group. Measured per draft metadata build on the host: 6.76 us -> 1.86 us. Co-Authored-By: Claude Fable 5.1 Signed-off-by: gf239 --- vllm/v1/worker/gpu/attn_utils.py | 9 ++++++++- 1 file changed, 8 insertions(+), 1 deletion(-) diff --git a/vllm/v1/worker/gpu/attn_utils.py b/vllm/v1/worker/gpu/attn_utils.py index 05949c25ac21..506b07c77578 100644 --- a/vllm/v1/worker/gpu/attn_utils.py +++ b/vllm/v1/worker/gpu/attn_utils.py @@ -288,6 +288,13 @@ def build_attn_metadata( attn_metadata: dict[str, Any] = {} num_kv_cache_groups = len(kv_cache_config.kv_cache_groups) for i in range(num_kv_cache_groups): + groups = attn_groups[i] + if not groups: + # This model owns no layer in this KV cache group. A drafter builds + # its attention groups from active_layer_names only, so + # init_attn_backend appends [] for every group it does not own, and + # everything below would be constructed and immediately discarded. + continue block_table = block_tables[i] slot_mapping = slot_mappings[i] # Per-group causal for hybrid drafters (mixed SWA/full attention). @@ -325,7 +332,7 @@ def build_attn_metadata( **common_attn_metadata_extra_kwargs, ) - for attn_group in attn_groups[i]: + for attn_group in groups: attn_metadata_builder = attn_group.get_metadata_builder(ubatch_idx) if for_cudagraph_capture: metadata = attn_metadata_builder.build_for_cudagraph_capture(