diff --git a/vllm/model_executor/layers/activation.py b/vllm/model_executor/layers/activation.py index ab4be153..958ec540 100644 --- a/vllm/model_executor/layers/activation.py +++ b/vllm/model_executor/layers/activation.py @@ -97,7 +97,22 @@ class GeluAndMul(CustomOp): d = x.shape[-1] // 2 output_shape = (x.shape[:-1] + (d, )) - out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + # ═══════════════════════════════════════════════════════════════ + # CCCL dispatch_transform.cuh CacheAsyncConfiguration pattern: + # Output tensor shape is deterministic from input shape. + # During decode, shapes are stable → cache to avoid cudaMalloc. + # CCCL: "This computation MUST NOT depend on runtime state ... + # since the result will be cached." + # ═══════════════════════════════════════════════════════════════ + _cache_key = (output_shape, x.dtype, x.device) + _cached = getattr(self, '_out_cache', {}).get(_cache_key) + if _cached is not None and _cached.shape == output_shape: + out = _cached + else: + out = torch.empty(output_shape, dtype=x.dtype, device=x.device) + if not hasattr(self, '_out_cache'): + self._out_cache = {} + self._out_cache[_cache_key] = out if self.approximate == "none": ops.gelu_and_mul(out, x) elif self.approximate == "tanh": diff --git a/vllm/model_executor/layers/layernorm.py b/vllm/model_executor/layers/layernorm.py index eb5be2d1..c2d8937d 100644 --- a/vllm/model_executor/layers/layernorm.py +++ b/vllm/model_executor/layers/layernorm.py @@ -85,7 +85,28 @@ class RMSNorm(CustomOp): residual_alpha, ) return x, residual - out = torch.empty_like(x) + # ═══════════════════════════════════════════════════════════════ + # CCCL dispatch_transform.cuh CacheAsyncConfiguration pattern: + # Element-wise transforms have deterministic output shapes. + # During decode, input shape is stable (num_seqs × hidden_dim). + # Cache the output tensor to avoid cudaMalloc on every step. + # + # CCCL: "This computation MUST NOT depend on runtime state ... + # since the result will be cached." + # + # RMSNorm is called 64× per forward pass (Qwen3.6 has 64 layers). + # Each call was doing torch.empty_like → cudaMalloc. + # With caching: 64 cudaMalloc calls → 0 per decode step. + # ═══════════════════════════════════════════════════════════════ + _cache_key = (x.shape, x.dtype, x.device) + _cached = getattr(self, '_out_cache', {}).get(_cache_key) + if _cached is not None and _cached.shape == x.shape: + out = _cached + else: + out = torch.empty_like(x) + if not hasattr(self, '_out_cache'): + self._out_cache = {} + self._out_cache[_cache_key] = out ops.rms_norm( out, x,